Skip to content

gh-131798: JIT: Propagate the result in _BINARY_OP_SUBSCR_TUPLE_INT #133003

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 4 commits into from
Apr 26, 2025

Conversation

tomasr8
Copy link
Member

@tomasr8 tomasr8 commented Apr 26, 2025

More context: #132851 (comment)

This propagates the information about tuple elements after _BINARY_OP_SUBSCR_TUPLE_INT when the RHS is a constant.
For example in

foo = (1, 2)
x = foo[0]  # _BINARY_OP_SUBSCR_TUPLE_INT

we can now deduce that x is 1.

Comment on lines 378 to 379
assert(index >= 0);
assert(index < sym_tuple_length(left));
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm not sure these asserts are actually needed. The instruction will deopt when this is not true so perhaps we can remove them?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Interesting, we are actually hitting the assertion:

Python/optimizer_cases.c.h:630: optimize_uops: Assertion `index < sym_tuple_length(left)' failed.

I guess I need to move the DEOPT_IF checks into a separate guard for it to work?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Check out the code for sym_tuple_length. It can return -1 if the length is not known. So for example if you propagate a PyTuple_Type, but not the length, it will fail.

You need to check that the length is not -1.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Got it, thanks! I somehow didn't realize you can know something is a tuple but not know its length..

long index = PyLong_AsLong(sym_get_const(ctx, right));
assert(index >= 0);
assert(index < sym_tuple_length(left));
res = sym_tuple_getitem(ctx, left, index);
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is a bit unfortunate. I would wish we could automatically evaluate this, but it seems tuples are a special category so we can't.

#132733

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I haven't fully digested your PR yet, but maybe there is some way to extend it to support tuples as well?

@Fidget-Spinner Fidget-Spinner merged commit 5e96e4f into python:main Apr 26, 2025
62 checks passed
@Fidget-Spinner
Copy link
Member

Thanks!

@tomasr8 tomasr8 deleted the jit-subscr-tuple branch April 26, 2025 19:15
Comment on lines +375 to +391
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);
}
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@tomasr8, this can be improved to handle abstract tuples that may not be constant (and it cleans things up, since it includes a length check):

Suggested change
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);
}
long index = PyLong_AsLong(sym_get_const(ctx,
sub_st));
// ...PyLong_AsLong can fail, need to check for error here...
res = sym_tuple_getitem(ctx, tuple_st, index)

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this can be improved to handle abstract tuples that may not be constant

@brandtbucher, I think this already handles abstract tuples? Assuming that abstract means those that have JIT_SYM_TUPLE_TAG rather than JIT_SYM_KNOWN_VALUE_TAG? For tuples that just have JIT_SYM_KNOWN_CLASS_TAG we don't know the length so there's nothing we can do I believe.

But yeah, we can simplify the code since _Py_uop_sym_tuple_getitem can handle unknown length. The only difference is that it sets the type to unknown rather than non null as we do here. Is that a problem? Shouldn't _Py_uop_sym_tuple_getitem always return at least JIT_SYM_NON_NULL_TAG?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Your current code is gated on sym_is_const, which abstract tuples fail. The rest of your understanding is correct, though.

Is that a problem? Shouldn't _Py_uop_sym_tuple_getitem always return at least JIT_SYM_NON_NULL_TAG?

Sounds like you found your next PR. :)

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ah, I completely misunderstood, sorry. I thought you were checking whether the tuple was const, not the index. Disregard!

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

no worries :) I think we can still simplify the code a bit though!

Comment on lines +380 to +383
if (tuple_length == -1) {
// Unknown length
res = sym_new_not_null(ctx);
}
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@brandtbucher
Following the previous discussion, I'm wondering if it is correct to return sym_new_not_null here rather than sym_new_unknown?

We have an abstract tuple and we don't know its length so it is impossible to be sure that the index is not out of bounds for the tuple. Given that, I think it is incorrect to return sym_new_not_null here as regular Python code could raise an IndexError. Should we change this to sym_new_unknown?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants