Better constant narrowing in the JIT optimizer

[brandtbucher]

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 (==, !=)
• _COMPARE_OP_INT + _GUARD_IS_*_POP (==, !=)
• _COMPARE_OP_STR + _GUARD_IS_*_POP (==, !=)
• _CONTAINS_OP + _GUARD_IS_*_POP (in, not in)
• _IS_OP + _GUARD_IS_*_POP (is, is not)

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 (==, !=, <, >, <=, >=)
• _COMPARE_OP_INT (==, !=, <, >, <=, >=) gh-131489
• _COMPARE_OP_STR (==, !=, <, >, <=, >=)
• _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

[Fidget-Spinner]

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]

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.

_GUARD_IS_FALSE_CHECK only sees a boolean, and doesn't "know" how it was created or what it represents.

[brandtbucher]

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]

I think we should sym_set_const(value, Py_False) in _GUARD_IS_FALSE_POP instead.

[brandtbucher]

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]

Talked to Brandt in private. Sadly, have to agree here that we either need to peek or need a new symbol.

[markshannon]

What we could do for narrowing type based on boolean guards is this:

• Add a new boolean type of symbol, with a a field for its input(s).
• Add a function sym_to_bool to convert symbols into bools and track the input.
• Add another function sym_guarded_boolean(sym, true/false) that narrows the symbol.

That shouldn't be that much code, and the optimizer code for TO_BOOL and GUARD_... is then almost trivial.

[brandtbucher]

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.

[tomasr8]

Hi @brandtbucher! Are any of the tasks up for grabs still? I haven't worked on the JIT before, but I'd like to try :)

[picnixz]

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

[brandtbucher]

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]

Would you like to add a case to optimizer_bytecodes.c for _CONTAINS_OP_SET, which sets the output type to &PyBool_Type?

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.