LLVM/MIR optimizations to simplify the new RFC3058 ? desugaring

[scottmcm]

With #84767 implementing the new ? desugaring, there's no mir-opt that can simplify the double-match it emits, as the arm identity pass only does Result->Result, and thus cannot handle Result->ControlFlow->Result.

That's not a blocker for checkin, as the arm identity pass is disabled by default currently because of things like #78628 and it depends on the MIR inliner being enabled anyway, but it sure would be nice to do better here.

Zulip conversation with some more context: https://rust-lang.zulipchat.com/#narrow/stream/189540-t-compiler.2Fwg-mir-opt/topic/Suggestions.20for.20a.20beginner/near/237739360

[Moxinilian]

I am interested in taking this one as a first contribution. I will probably have time to work on it this week.

[Moxinilian]

For interested people, conversation on Zulip moved to https://rust-lang.zulipchat.com/#narrow/stream/189540-t-compiler.2Fwg-mir-opt/topic/Optimizing.20a.20match.20on.20a.20just-constructed.20enum.20.2385133

[michaelmaitland]

I see this issue marked as open. Is it being worked on or available for me to pickup?

[Moxinilian]

It is being worked on, see #85646. Any suggestion/help is welcome, though!

[BGR360]

Is this issue now resolved after @Moxinilian 's work?

[Moxinilian]

Right now it is not. Right now this is blocked by LLVM doing a sub-optimal transformation that loses the optimization at a very early stage, preventing the duplicated tails to be merged back into one. I've tried investigating if we could not merge them back directly in MIR. While possible in theory, my understanding was that MIR still abstracts away too much details (notably on the discriminant aspect) for the algorithm to be reasonably complex. So right now we either need a breakthrough on that front or work in the LLVM codebase. I want to start investigating more in the LLVM side of things, but I have not successfully allocated enough time for it as of now.

[Moxinilian]

More on simplifying in MIR:

What LLVM is struggling with is keeping the fact that variant switches are exhaustive, so it can merge the blocks that do the same thing except for different constants (representing the variant IDs). What we could do is do that merging directly in MIR, in a pass close to the end of the pipeline. A naive approach would be to take a switch-terminated block bb that ends in (semantically)

if x == 1: goto bb1
else if x == 2: goto bb2
else: unreachable

Then, replace any instance of the constant 1 in bb1 with x before x is reassigned and every instance of 2 in bb2 with x before x is reassigned. Now, if the new versions of bb1 and bb2 are equal, the blocks can safely be merged. That seems to be sound to me but it would not be very useful here because:

• Maybe other optimizations transformed the blocks asymmetrically in such a way that they are semantically equivalent but not syntactically equal (for example if an useless affectation got inserted for some reason). One would need to make a canonicalization pass to fix this, which probably could resolve some of those issues. However, the second point is the most important:
• Discriminants must be constant in MIR form, so we can't replace them with an use of x, even though later down the pipeline it would be lowered into that. This kills pretty much all of our use cases for this issue.

Note that if MIR was in SSA form, this would probably be more effective in practice because x would never be re-defined.

[Moxinilian]

More on fixing LLVM:

Right now, rustc calls into LLVM with a pass pipeline that starts with the SimplifyCFG LLVM pass that uses the LowerSwitch utility to turn switches into lower level constructs. An issue with that pass is that it is opinionated in what is worth simplifying away, which includes unreachable switch targets. So for example, the following LLVM IR:

define i64 @demo(i64 %x) {
entry:
    switch i64 %x, label %bb3 [
        i64 0, label %bb1
        i64 1, label %bb2
    ]
bb1:
    ret i64 0
bb2:
    %0 = icmp eq i64 %x, 100 ; this will necessarily be false because %x == 1
    br i1 %0, label %bb4, label %bb5
bb3:
    unreachable
bb4:
    ret i64 200
bb5:
    ret i64 10
}

is turned into

define i64 @demo(i64 %x) {
entry:
    %switch = icmp ult i64 %x, 1
    %0 = icmp eq i64 %x, 100
    %. = select i1 %0, i64 200, i64 10
    %merge = select i1 %switch, i64 0, i64 %.
    ret i64 %merge
}

which throws away the exhaustive nature of the switch, making in this example further passes not able to optimize the useless comparison away.

The code of LowerSwitch seems to contain an explicit part where it removes unreachable branches in switches. Forcing this off might change this behavior and make it work as intended. I have not gotten to try that yet, but plan to.

[BGR360]

Oh, that's annoying. But very intriguing. Thanks for the explanation!

[LebedevRI]

FWIW this has been fixed in llvm trunk in llvm/llvm-project@371fcb7.

[nikic]

So ... what is the actual Rust code we're supposed to be optimizing here? Is the goal to get rid of the icmp + zext in https://rust.godbolt.org/z/nd9fEx3ed or is it something else? (The referenced LLVM change does not get rid of those.)

[scottmcm]

@nikic That's right. try { x? } not being a nop, as in your example, is the easiest demonstration.

As I understand it, the icmp shows up because there's a trunc in there originally, but LLVM has lost track of the fact that the truncation doesn't ever wrap.

(My uninformed thought: It's odd to me that there's no nsw/nuw flags for trunc, as a way to say "I promise this came from a [sz]ext". If there were, then we could emit those flags in the discriminant loading, and avoid the icmp.)

[scottmcm]

Actually, here's a better cut-down version of the example. This is (essentially) the first part of ?, but with no trait calls involved to muddle things: https://rust.godbolt.org/z/MWxdTcGn6

use std::ops::ControlFlow;
pub fn demo(x: Result<u32, i32>) -> ControlFlow<i32, u32> {
    match x {
        Ok(v) => ControlFlow::Continue(v),
        Err(e) => ControlFlow::Break(e),
    }
}

Ok and Continue are both discriminant 0; Err and Break are both discriminant 1.

It looks like the core of what we emit there is

  %8 = load i32, i32* %7, align 4, !dbg !10, !range !11
  %_2 = zext i32 %8 to i64, !dbg !10
  switch i64 %_2, label %bb2 [
    i64 0, label %bb3
    i64 1, label %bb1
  ], !dbg !12

where that range metadata is

!11 = !{i32 0, i32 2}

So AFAIK that PR above was about trying to avoid turning the switch into an icmp+br, which lost that range info.

[nikic]

Okay, so I think the PhaseOrdering change is necessary to avoid folding the switch before SROA (of course, the alternative solution is to make SROA emit range assumptions). But it's not sufficient:

*** IR Dump After InstCombinePass on _ZN4test4test17h60e5e3f4ec590b13E ***
; Function Attrs: nonlazybind uwtable
define { i64, i64 } @_ZN4test4test17h60e5e3f4ec590b13E(i64 %x.0, i64 %x.1) unnamed_addr #0 {
start:
  %0 = call fastcc { i64, i64 } @"_ZN79_$LT$core..result..Result$LT$T$C$E$GT$$u20$as$u20$core..ops..try_trait..Try$GT$6branch17hed04f6fc65ae78f0E"(i64 %x.0, i64 %x.1)
  %.fca.0.extract1 = extractvalue { i64, i64 } %0, 0
  %.fca.1.extract2 = extractvalue { i64, i64 } %0, 1
  switch i64 %.fca.0.extract1, label %bb3 [
    i64 0, label %bb2
    i64 1, label %bb4
  ]

bb3:                                              ; preds = %start
  unreachable

bb2:                                              ; preds = %start
  br label %bb6

bb4:                                              ; preds = %start
  %1 = call fastcc i64 @"_ZN153_$LT$core..result..Result$LT$T$C$F$GT$$u20$as$u20$core..ops..try_trait..FromResidual$LT$core..result..Result$LT$core..convert..Infallible$C$E$GT$$GT$$GT$13from_residual17h7fc30301c010dde1E"(i64 %.fca.1.extract2)
  br label %bb6

bb6:                                              ; preds = %bb2, %bb4
  %.sroa.3.0 = phi i64 [ %1, %bb4 ], [ %.fca.1.extract2, %bb2 ]
  %.sroa.0.0 = phi i64 [ 1, %bb4 ], [ 0, %bb2 ]
  %2 = insertvalue { i64, i64 } undef, i64 %.sroa.0.0, 0
  %3 = insertvalue { i64, i64 } %2, i64 %.sroa.3.0, 1
  ret { i64, i64 } %3
}
*** IR Dump After SimplifyCFGPass on _ZN4test4test17h60e5e3f4ec590b13E ***
; Function Attrs: nonlazybind uwtable
define { i64, i64 } @_ZN4test4test17h60e5e3f4ec590b13E(i64 %x.0, i64 %x.1) unnamed_addr #0 {
start:
  %0 = call fastcc { i64, i64 } @"_ZN79_$LT$core..result..Result$LT$T$C$E$GT$$u20$as$u20$core..ops..try_trait..Try$GT$6branch17hed04f6fc65ae78f0E"(i64 %x.0, i64 %x.1)
  %.fca.0.extract1 = extractvalue { i64, i64 } %0, 0
  %.fca.1.extract2 = extractvalue { i64, i64 } %0, 1
  %switch = icmp ult i64 %.fca.0.extract1, 1
  br i1 %switch, label %bb6, label %bb4

bb4:                                              ; preds = %start
  %1 = call fastcc i64 @"_ZN153_$LT$core..result..Result$LT$T$C$F$GT$$u20$as$u20$core..ops..try_trait..FromResidual$LT$core..result..Result$LT$core..convert..Infallible$C$E$GT$$GT$$GT$13from_residual17h7fc30301c010dde1E"(i64 %.fca.1.extract2)
  br label %bb6

bb6:                                              ; preds = %start, %bb4
  %.sroa.3.0 = phi i64 [ %1, %bb4 ], [ %.fca.1.extract2, %start ]
  %.sroa.0.0 = phi i64 [ 1, %bb4 ], [ 0, %start ]
  %2 = insertvalue { i64, i64 } undef, i64 %.sroa.0.0, 0
  %3 = insertvalue { i64, i64 } %2, i64 %.sroa.3.0, 1
  ret { i64, i64 } %3
}

We'd need either InstCombine or SimplifyCFG to fold %.sroa.0.0 into %.fca.0.extract1.

I believe SimplifyCFG would do that as part of ForwardSwitchConditionToPHI, but this is only done at the very end of the pipeline.

[nikic]

Just remembered that we have https://github.com/llvm/llvm-project/blob/fdfe26ddbeb1a5eb6106a6a27d6f8240deca543f/llvm/lib/Transforms/InstCombine/InstCombinePHI.cpp#L1256 in InstCombine. I guess it would be natural to extend to that to switches and not just conditional branches.

[nikic]

Upstream fix is llvm/llvm-project@4010a7a.

[nikic]

The LLVM side of this is fixed now, and #100693 added a codegen test. Is a MIR transform still needed?

[scottmcm]

I'll close this one because it does look like the match-of-match, as ? does, is now working great.

Unfortunately, however, when just matching the result directly and rebuilding, it seems to only work for some widths. The codegen test works with Result<i32, u32>, but doesn't optimize away any more for Result<i64, u64>. I've opened #101210 to track that separately, @nikic .