Conditional backward edges should help "warm up" code

[brandtbucher]

#93229 introduced a regression in how aggressively we quicken some for loops. Minimal example:

def f(x: bool) -> None:
    for i in range(1_000_000):
        if x:
            pass

f(True) will quicken this code, but f(False) will not, even though both contain the same number of back edges. The issue is that we only quicken on unconditional backwards jumps, not on conditional ones.

We've known about this limitation for some time, in particular with regard to while loops. Since we check the loop condition at the bottom of while loops, one call is not enough to quicken w:

def w() -> None:
    i = 0
    while i < 1_000_000:
        i += 1

@markshannon has expressed a preference for having all branches be forward (i.e. replacing backward POP_JUMP_IF_FALSE(x) instructions with POP_JUMP_FORWARD_IF_TRUE(1); JUMP_BACKWARD(x) in the assembler). @iritkatriel believes that this shouldn't be too difficult, based on recent assembler rewrites.

CC @sweeneyde

Linked PRs

• gh-93554: add test for quickening of code in loops ending with conditional statement #119485

[sweeneyde]

Darn. Switching over to forward-only conditional jumps makes sense to me then.

[sweeneyde]

If it's easier to just revert #93229 for now, that is okay as well.

[iritkatriel]

@brandtbucher the assembler adds a basicblock with a branch here .

[brandtbucher]

This would need to happen later though, right? Like, at the last possible moment, after we've ordered the blocks and we're actually emitting concrete jump instructions?

The bummer is that I don't see a way to make this play nicely with the peepholer's jump elimination, unless we add a separate jump-to-jump elimination pass there. We may just need to accept that some optimization opportunities will slip through the cracks.

[sweeneyde]

The bummer is that I don't see a way to make this play nicely with the peepholer's jump elimination, unless we add a separate jump-to-jump elimination pass there.

Couldn't we do jump-to-jump-elimination first, then afterward convert conditional backward jumps to (conditional forward jump + JUMP_BACKWARD)? I don't know how common chains of three jumps in a row are, but at least then those would be converted to only two jumps in a row (conditional_forward+backward).

[brandtbucher]

Most jump elimination will still work fine, but I'm anticipating weird edge-cases with this new transformation.

I'm picturing something like POP_JUMP_IF_FALSE(<backward>); JUMP(<forward>). Peepholing can't do anything with that, but the assembler will turn it into POP_JUMP_FORWARD_IF_TRUE(1); JUMP_BACKWARD(<backward>); JUMP_FORWARD(<forward>). So we miss an opportunity to eliminate a jump (and potentially a whole block).

[brandtbucher]

It's related to the issue with FOR_ITER jumping to JUMP now. By the time we find out if it's a legal transformation, it's too late.

[sweeneyde]

Would that be solved by having separate cases like this?

case [COND_JUMP(backward to A), JUMP(forward to B)]:
    replace with [INVERTED_COND_JUMP(forward to B), JUMP_BACKWARD(to A)]
case [COND_JUMP(backward to A), SOMETHING_ELSE]:
    replace with [INVERTED_COND_JUMP(forward to SOMETHING_ELSE), JUMP_BACKWARD(to A), SOMETHING_ELSE]

In the first case, Jump-To-Jump elimination will have already folded JUMP(forward to B) as much as possible, so there should be nothing left to fold.

[iritkatriel]

Can we move jump elimination later, just before we emit the code or is there a reason it needs to be earlier?

[sweeneyde]

Can we move jump elimination later, just before we emit the code or is there a reason it needs to be earlier?

From what I can gather, jump-threading can mess up directions of conditional jumps, but (assuming we have special logic for conditional-followed-by-unconditional like my other comment) inverting conditional jumps can't introduce any new opportunities for jump-threading.

I'm imagining a situation where the front-end makes a monster like this:

    JUMP to C
A:
    JUMP to Y
B:
    JUMP to X
C:
    [...]
    COND_JUMP to Z
elseblock:
    [...]
    RETURN
X:
    JUMP to C
Y:
    JUMP to B
Z:
    JUMP to A

Conditional-Jump-Inversion, then Jump-Threading

If we attempt to do conditional jump inversion first, there's nothing to do: the conditional jump is already forward. Then when we go to do jump threading, we get

Jump-Inversion (nothing to do), followed by Jump-Threading

    JUMP to C
A:
    JUMP to C
B:
    JUMP to C
C:
    [...]
    COND_JUMP to C
elseblock:
    [...]
    RETURN
X:
    JUMP to C
Y:
    JUMP to C
Z:
    JUMP to C

Then we're left with a backward conditional jump, which is the issue.

I suppose one alternative would be to limit the power of jump-threading so that forward jumps always remain forward jumps. It would probably be safe to point the conditional jump at the last instruction in the chain that is forward of it instead (X in this case). That might be okay, but I think we might then need a different algorithm for jump threading, maybe using Floyd's tortoise/hare or something.

Or maybe such long messy chains are so rare that it wouldn't be too bad to just thread jumps as long as things are forward, and then give up and accept any remaining jumps to jumps.

Jump-Threading, then Conditional-Jump-Inversion

After Jump-Threading first (same as today):

    JUMP to C
A:
    JUMP to C
B:
    JUMP to C
C:
    [...]
    COND_JUMP to C
elseblock:
    [...]
    RETURN
X:
    JUMP to C
Y:
    JUMP to C
Z:
    JUMP to C

Then Conditional-Jump-Inversion afterward:

    JUMP to C
A:
    JUMP to C
B:
    JUMP to C
C:
    [...]
    INVERTED_COND_JUMP to elseblock
    COND_JUMP_TO_C
elseblock:
    [...]
    RETURN
X:
    JUMP to C
Y:
    JUMP to C
Z:
    JUMP to C

(In either case, we can then remove unreachable blocks and no-op jumps.)

One other thing I'd note is that we currently don't fold JUMP -> CONDIONAL_JUMP, but if we did, that would create another way for jump-threading to mess up the direction of conditional jump.

Apologies for the wall of text 😅