Rust insufficiently optimizes loop { match { } } state machines

[tuxmark5]

This is how a typical state machine might look in Rust:

#[inline(never)]
fn print(s: &str) {
    println!("{}", s);
}

#[inline(never)]
fn exec(mut s: i32) {
    loop {
        match s {
            0 => { print("0"); s = 2; },
            1 => { print("1"); s = 3; },
            2 => { print("2"); s = 1; },
            _ => return
        }
    }
}

fn main() {
    exec(0);
}

One would expect that the first jump from function start into one of the states (0, 1, 2 or _) would be compiled as a jumptable or a series of cmp/conditional jump instructions and once the execution is in one of the states the remaining jumps would be deterministic (non-conditional) jumps.

And this is how equivalent code in gcc is optimized.

However Rust/LLVM at the end of each state performs needless comparisons to determine to which label to jump to next, even though the jump target can be computed statically. You can inspect the generated assembly in Rust Playground.

This is probably more of a LLVM issue rather than Rust directly, but it's important to get this fixed/optimized, because now it's impossible to efficiently write such state machines in Rust (because there is no goto)

Meta

Tested with all versions available in Rust playground:

• 1.49.0
• 1.50.0-beta.2 (2020-12-31 25b3db3)
• 1.51.0-nightly (2020-12-30 e226704)

[the8472]

One would expect that the first jump from function start into one of the states (0, 1, 2 or _) would be compiled as a jumptable or a series of cmp/conditional jump instructions

Atleast that part happens once you add more states.

[nagisa]

Didn't look too deeply into it but the IR fed to LLVM around the loop-match construct is not that different between rustc and clang. However with clang it becomes a switchtable at llvm-ir level, and we rely on the codegen backend to do so with rust https://godbolt.org/z/9Erv94

FWIW from what I can tell the additional comparisons can be blamed on s = <new> state setting and the phi that results from it.

[oli-obk]

I thought we already had an optimization for this

    bb1: {
        switchInt(_1) -> [0_i32: bb3, 1_i32: bb5, 2_i32: bb7, otherwise: bb2]; // scope 0 at src/lib.rs:10:13: 10:14
    }
    bb4: {
        StorageDead(_3);                 // scope 0 at src/lib.rs:10:29: 10:30
        StorageDead(_4);                 // scope 0 at src/lib.rs:10:30: 10:31
        StorageDead(_2);                 // scope 0 at src/lib.rs:10:30: 10:31
        _1 = const 2_i32;                // scope 0 at src/lib.rs:10:32: 10:37
        goto -> bb1;                     // scope 0 at src/lib.rs:9:9: 14:10
    }

just looking locally at this, the goto -> bb1 should become goto -> bb7. Basically if a block assigns a constant to a local and then jumps to a block which does not touch that local and ends up conditionally branching on the local, then we can copy all statements from the branching block to the current block and make it an unconditional jump to the right target.

[simonvandel]

AFAICT this is optimized with #80475

[nikic]

LLVM can optimize this with the -C llvm-args=-jump-threading-across-loop-headers option. By default LLVM does not perform jump threading across loop headers, because it cannot reliably determine profitability for this case right now. Doing this naively results in natural loops becoming irreducible, which will pessimize all following loop optimizations.

[camelid]

AFAICT this is optimized with #80475

That PR is merged now. Did that fix this issue?

[tuxmark5]

Testing it with (2021-03-20 61edfd5) on the playground shows no improvement.

[oli-obk]

Testing it with (2021-03-20 61edfd5) on the playground shows no improvement.

The playground does not actually run mir optimizations of level 3 or above, so you would need to run this locally with nightly and use -Zmir-opt-level=3

Doing this naively results in natural loops becoming irreducible, which will pessimize all following loop optimizations.

This is worrysome. Our optimization as written won't do anything fancy, so it will make loops more complex. The perf results on that optimization are mostly in the noise (ignoring the CTFE stress tests), though without a dedicated compile-time and runtime test, it's hard to tell. Considering that the compiler bootstrap did not regress, I don't think the opt messes with anything too badly.

[camelid]

Testing it with (2021-03-20 61edfd5) on the playground shows no improvement.

The playground does not actually run mir optimizations of level 3 or above, so you would need to run this locally with nightly and use -Zmir-opt-level=3

Or you can use Godbolt: https://rust.godbolt.org/z/doK6W9jPT

(I don't what the Assembly is supposed to look like so I can't tell if this is fixed or not.)

[tuxmark5]

Unfortunately it still doesn't seem fixed. Assembly from godbolt link:

example::exec:  << The start of exec() function in rust
        sub     rsp, 8
        mov     dword ptr [rsp + 4], edi
.LBB2_1: << The block that implements the 'switch' 
        mov     eax, dword ptr [rsp + 4]
        mov     dword ptr [rsp], eax
        test    eax, eax
        je      .LBB2_3
        jmp     .LBB2_9
.LBB2_9: 
        mov     eax, dword ptr [rsp]
        sub     eax, 1
        je      .LBB2_5
        jmp     .LBB2_10
.LBB2_10:
        mov     eax, dword ptr [rsp]
        sub     eax, 2
        je      .LBB2_7
        jmp     .LBB2_2
.LBB2_2:
        pop     rax
        ret
.LBB2_3: << case 0
        lea     rdi, [rip + .L__unnamed_2]
        mov     esi, 1
        call    qword ptr [rip + example::print@GOTPCREL]
        mov     dword ptr [rsp + 4], 2 << explicit store to stack: s = 2
        jmp     .LBB2_1 << BAD: jumps back to the 'switch', but s value is known
.LBB2_5: << case 1
        lea     rdi, [rip + .L__unnamed_3]
        mov     esi, 1
        call    qword ptr [rip + example::print@GOTPCREL]
        mov     dword ptr [rsp + 4], 3 << explicit store to stack: s = 3
        jmp     .LBB2_1 << BAD: jumps back to the 'switch' , but s value is known
.LBB2_7: << case 2
        lea     rdi, [rip + .L__unnamed_4]
        mov     esi, 1
        call    qword ptr [rip + example::print@GOTPCREL]
        mov     dword ptr [rsp + 4], 1 << explicit store to stack: s = 1
        jmp     .LBB2_1 << BAD: jumps back to the 'switch' , but s value is known

I've annotated it a bit with << to make it easier to understand. Basically all the match arms eventually jump back to the 'switch' block, where s is compared/tested.

[oli-obk]

Maybe LLVM "de-optimizes" our MIR again. Not sure, we'll have to look at the MIR before LLVM, at LLVM-IR without opts and at LLVM-IR with opts to be sure. I'm not sure how to coax godbolt into doing MIR dumps though.