gh-131798: JIT: Propagate the result in _BINARY_OP_SUBSCR_TUPLE_INT (GH-133003)

This commit is contained in:
Tomas R. 2025-04-26 20:47:55 +02:00 committed by GitHub
parent 56c88e4e8d
commit 5e96e4fca8
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 61 additions and 1 deletions

View file

@ -1923,6 +1923,23 @@ class TestUopsOptimization(unittest.TestCase):
self.assertNotIn("_GUARD_TOS_INT", uops)
self.assertIn("_CALL_LEN", uops)
def test_binary_op_subscr_tuple_int(self):
def testfunc(n):
x = 0
for _ in range(n):
y = (1, 2)
if y[0] == 1: # _COMPARE_OP_INT + _GUARD_IS_TRUE_POP are removed
x += 1
return x
res, ex = self._run_with_optimizer(testfunc, TIER2_THRESHOLD)
self.assertEqual(res, TIER2_THRESHOLD)
self.assertIsNotNone(ex)
uops = get_opnames(ex)
self.assertIn("_BINARY_OP_SUBSCR_TUPLE_INT", uops)
self.assertNotIn("_COMPARE_OP_INT", uops)
self.assertNotIn("_GUARD_IS_TRUE_POP", uops)
def global_identity(x):
return x

View file

@ -0,0 +1,2 @@
Propagate the return type of ``_BINARY_OP_SUBSCR_TUPLE_INT`` in JIT. Patch
by Tomas Roun

View file

@ -370,6 +370,27 @@ dummy_func(void) {
res = sym_new_type(ctx, &PyUnicode_Type);
}
op(_BINARY_OP_SUBSCR_TUPLE_INT, (tuple_st, sub_st -- res)) {
assert(sym_matches_type(tuple_st, &PyTuple_Type));
if (sym_is_const(ctx, sub_st)) {
assert(PyLong_CheckExact(sym_get_const(ctx, sub_st)));
long index = PyLong_AsLong(sym_get_const(ctx, sub_st));
assert(index >= 0);
int tuple_length = sym_tuple_length(tuple_st);
if (tuple_length == -1) {
// Unknown length
res = sym_new_not_null(ctx);
}
else {
assert(index < tuple_length);
res = sym_tuple_getitem(ctx, tuple_st, index);
}
}
else {
res = sym_new_not_null(ctx);
}
}
op(_TO_BOOL, (value -- res)) {
int already_bool = optimize_to_bool(this_instr, ctx, value, &res);
if (!already_bool) {

View file

@ -617,8 +617,28 @@
}
case _BINARY_OP_SUBSCR_TUPLE_INT: {
JitOptSymbol *sub_st;
JitOptSymbol *tuple_st;
JitOptSymbol *res;
res = sym_new_not_null(ctx);
sub_st = stack_pointer[-1];
tuple_st = stack_pointer[-2];
assert(sym_matches_type(tuple_st, &PyTuple_Type));
if (sym_is_const(ctx, sub_st)) {
assert(PyLong_CheckExact(sym_get_const(ctx, sub_st)));
long index = PyLong_AsLong(sym_get_const(ctx, sub_st));
assert(index >= 0);
int tuple_length = sym_tuple_length(tuple_st);
if (tuple_length == -1) {
res = sym_new_not_null(ctx);
}
else {
assert(index < tuple_length);
res = sym_tuple_getitem(ctx, tuple_st, index);
}
}
else {
res = sym_new_not_null(ctx);
}
stack_pointer[-2] = res;
stack_pointer += -1;
assert(WITHIN_STACK_BOUNDS());