Repository navigation
Better constant narrowing in the JIT optimizer #130415
Description
Please don't work on this. I'm planning on sprinting on this with new contributors at an event this weekend.
Currently, the JIT optimizer uses the _COMPARE_OP family, _CONTAINS_OP, _IS_OP, and the _TO_BOOL family to narrow the types of the input values and the type of the output value.
However, by "peeking ahead" and seeing how a value will be used, we can narrow these types to constants as well. As a simple example, consider _TO_BOOL_INT + _GUARD_IS_FALSE_CHECK on an unknown value. After the _TO_BOOL_INT, it can be narrowed to a known class, int (we do this today). However, after the _GUARD_IS_FALSE_CHECK, we can actually narrow it to a constant value, 0.
An example implementation of this idea for _TO_BOOL_BOOL is here: main...brandtbucher:cpython:hack-night-to-bool-bool
I've divided this work into 3 "waves" of increasing complexity. Tasks in bold are probably a bit harder, tasks in italics are probably a bit easier.
Narrow types to constants in branches involving truthiness:
- _TO_BOOL + _GUARD_IS_*_POP gh-130659
- _TO_BOOL_BOOL + _GUARD_IS_*_POP gh-130659
- _TO_BOOL_INT + _GUARD_IS_*_POP gh-130772
- _TO_BOOL_LIST + _GUARD_IS_*_POP
- _TO_BOOL_STR + _GUARD_IS_*_POP gh-130476
Narrow types to constants in branches involving comparisons with a constant:
- _COMPARE_OP + _GUARD_IS_*_POP (==, !=)
- _COMPARE_OP_FLOAT + _GUARD_IS_*_POP (==, !=) gh-130415: Narrow types to constants in branches involving specialized comparisons with a constant #144150
- _COMPARE_OP_INT + _GUARD_IS_*_POP (==, !=) gh-130415: Narrow types to constants in branches involving specialized comparisons with a constant #144150
- _COMPARE_OP_STR + _GUARD_IS_*_POP (==, !=) gh-130415: Narrow types to constants in branches involving specialized comparisons with a constant #144150
- _CONTAINS_OP + _GUARD_IS_*_POP (in, not in)
- _IS_OP + _GUARD_IS_*_POP (is, is not) gh-130415: Narrowing to constants in branches involving is comparisons with a constant #143895
Evaluate comparisons involving two constants:
This is related, but a bit more involved, since we need a way to pop two values from the stack and push a constant (_POP_TWO_LOAD_CONST_INLINE_BORROW). We should also teach remove_unneeded_uops about this new instruction.
- _COMPARE_OP (==, !=, <, >, <=, >=)
- _COMPARE_OP_FLOAT (==, !=, <, >, <=, >=) gh-137062
- _COMPARE_OP_INT (==, !=, <, >, <=, >=) gh-131489
- _COMPARE_OP_STR (==, !=, <, >, <=, >=) gh-137062
- _CONTAINS_OP (in, not in)
- _IS_OP (is, is not)
Linked PRs
- GH-130415: Narrow str to "" based on boolean tests #130476
- GH-130415: Optimize JIT path for _TO_BOOL_INT branching #130477
- GH-130415: Use boolean guards to narrow types to values in the JIT #130659
- GH-130415: Narrow int to 0 based on boolean tests #130772
- GH-130415: Optimize constant comparison in JIT builds #131489
- GH-130415: Remove redundant sym_matches_type calls in the JIT optimizer #131778
- GH-130415: Improve the JIT's unneeded uop removal pass #132333
- gh-130415: Eliminate guards for constant CALL_BUILTIN_O/FAST #132708
- GH-130415: Use POP_TWO_LOAD_CONST_INLINE_BORROW #134241
- GH-132732: Use pure op machinery to optimize COMPARE_OP_INT/FLOAT/STR #137062
- gh-130415: Narrowing to constants in branches involving is comparisons with a constant #143895
- gh-130415: Narrow types to constants in branches involving specialized comparisons with a constant #144150
Activity
Fidget-Spinner commented on Feb 22, 2025
Wait I'm a bit confused, why can't we narrow it at the point of the guard, instead of looking ahead by 1? I don't generally like peeking ahead opcodes as it breaks the encapsulation of each opcode in the DSL.
brandtbucher commented on Feb 22, 2025
_GUARD_IS_FALSE_CHECK only sees a boolean, and doesn't "know" how it was created or what it represents.
brandtbucher commented on Feb 22, 2025
Another option could be to have new symbolic types like JitBoolOf(JitOptSymbol *source), JitComparison(JitOptSymbol *lhs, int op, int arg, JitOptSymbol *rhs). But correctly resolving those introduces a lot of complexity, and is probably best left as a future improvement if we deem it's worth it. This is a bit more pragmatic, and handles all of the "obvious" cases quite easily.
At the very least, this gets something simple working and adds some tests for the desired behavior that we can improve the implementation of later.
Fidget-Spinner commented on Feb 22, 2025
I think we should sym_set_const(value, Py_False) in _GUARD_IS_FALSE_POP instead.
brandtbucher commented on Feb 22, 2025
I think you might be focusing too much on the bool example. For most ops, we've already transformed the original value into an abstract bool, and don't have it anymore.
Your suggestion wouldn't work for _TO_BOOL_INT + _GUARD_IS_FALSE_POP, or _COMPARE_OP_FLOAT + _GUARD_IS_TRUE_POP, etc. The "value" seen by _GUARD_IS_FALSE_POP is different from the "value" seen by _TO_BOOL_INT or _COMPARE_OP_FLOAT. It's been converted to a bool by that point, and we've lost the original input(s).
Fidget-Spinner commented on Feb 22, 2025
Talked to Brandt in private. Sadly, have to agree here that we either need to peek or need a new symbol.
markshannon commented on Feb 24, 2025
What we could do for narrowing type based on boolean guards is this:
That shouldn't be that much code, and the optimizer code for TO_BOOL and GUARD_... is then almost trivial.
brandtbucher commented on Feb 24, 2025
Yeah, after talking to @markshannon we decided that I'll get a PR up with the necessary infrastructure, and we'll update the PRs to use it once it's in.
6 remaining items
picnixz commented on Mar 30, 2025
@tomasr8 I've updated the issue with the related PRs and what they were solving so you have enough tasks to pick up I think :')
pythonGH-130415: Remove redundant sym_matches_type calls in the JIT o…
brandtbucher commented on Apr 3, 2025
Thanks for your interest, @tomasr8! I think we actually may have some good issues over on GH-131798.
Would you like to add a case to optimizer_bytecodes.c for _CONTAINS_OP_SET, which sets the output type to &PyBool_Type? That way, the JIT knows that code like foo in some_set returns a bool. It should look just like the existing handler for _CONTAINS_OP. (Don't forget to run make regen-cases!)
This allows the JIT to automatically remove an unnecessary _TO_BOOL_BOOL instruction when people do condition = foo in some_set; if condition: .... The test for this should probably look something like test_compare_pop_two_load_const_inline_borrow, except you'll need to write a new testfunc and assert that _CONTAINS_OP_SET is in the uops, but _TO_BOOL_BOOL isn't. The latter check should fail before your change, and pass after.
Let me know if you start working on it, and feel free to ping me anywhere (here, Discord, email) if you get stuck.
tomasr8 commented on Apr 3, 2025
Sounds good! I will take a look :) Thanks for the detailed explanation of the issue, that is super helpful and it's pretty clear to me now what I need to do.
GH-130415: Improve the JIT's unneeded uop removal pass (GH-132333)
Fidget-Spinner commented on May 19, 2025
I think we should pause this for now until #132733 is merged. Adding constant propagation code is error-prone and tedious. With that PR, contributors can just mark a bytecode as pure and get it for free for everything. Then we could just add more to the cases generator to improve on the constant propagation code, rather than writing it manually ourselves.
brandtbucher commented on Jun 6, 2025
That makes sense for the third category ("evaluate comparisons involving two constants"). The second category ("narrow types to constants in branches involving comparisons with a constant") is probably still worth doing, but will require some new symbolic values to represent equality and inequality (similar to the new truthiness symbols).
Fidget-Spinner commented on Jun 27, 2025
This PR is landed. So things are no longer blocked!
savannahostrowski commented on Sep 18, 2025
Just crossed out IS_OP since @brandtbucher and I had discussed this and decided it was too complicated to handle.
gh-130415: Narrowing to constants in branches involving `is` comparis…
gh-130415: Narrow types to constants in branches involving specialize…
Fidget-Spinner commented on Jan 24, 2026
_COMPARE_OP is a little difficult as it's an overloaded op, and can call arbitrary Python code.
_CONTAINS_OP as well.
Thanks to @reidenong , we completed IS_OP and COMPARE_OP_(TYPE). In the future, when we start recording types/values, we can narrow the LHS if its comparison is side-effect free #144179
This issue is done. Thanks everyone!
pythongh-130415: Narrowing to constants in branches involving `is` co…
pythongh-130415: Narrow types to constants in branches involving spec…