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
• _TO_BOOL_BOOL + _GUARD_IS_*_POP
• _TO_BOOL_INT + _GUARD_IS_*_POP
• _TO_BOOL_LIST + _GUARD_IS_*_POP
• _TO_BOOL_STR + _GUARD_IS_*_POP

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 (==, !=, <, >, <=, >=)
• _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

[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.