Large match statements that compile down to small assembly don't get inlined automatically

[pozix604]

It seems that function inlining heuristics fail in the following use case and does not inline when it should. Given huge match statement:

pub enum XXX {
    A1, B1, C1, D1, E1, F1, G1, H1, I1, J1, K1, L1, M1, N1, O1, P1, Q1, R1, S1, T1, U1, V1, W1, X1, Y1, Z1,
    A2, B2, C2, D2, E2, F2, G2, H2, I2, J2, K2, L2, M2, N2, O2, P2, Q2, R2, S2, T2, U2, V2, W2, X2, Y2, Z2,
}

impl std::convert::TryFrom<u8> for XXX {
    type Error = ();
    fn try_from(x: u8) -> Result<Self, Self::Error> {
            match () {
                    _ if x == XXX::A1 as u8 => Ok(XXX::A1),
                    _ if x == XXX::B1 as u8 => Ok(XXX::B1),
                    _ if x == XXX::C1 as u8 => Ok(XXX::C1),
                    _ if x == XXX::D1 as u8 => Ok(XXX::D1),
                    _ if x == XXX::E1 as u8 => Ok(XXX::E1),
                    _ if x == XXX::F1 as u8 => Ok(XXX::F1),
                    _ if x == XXX::G1 as u8 => Ok(XXX::G1),
                    _ if x == XXX::H1 as u8 => Ok(XXX::H1),
                    _ if x == XXX::I1 as u8 => Ok(XXX::I1),
                    _ if x == XXX::J1 as u8 => Ok(XXX::J1),
                    _ if x == XXX::K1 as u8 => Ok(XXX::K1),
                    _ if x == XXX::L1 as u8 => Ok(XXX::L1),
                    _ if x == XXX::M1 as u8 => Ok(XXX::M1),
                    _ if x == XXX::N1 as u8 => Ok(XXX::N1),
                    _ if x == XXX::O1 as u8 => Ok(XXX::O1),
                    _ if x == XXX::P1 as u8 => Ok(XXX::P1),
                    _ if x == XXX::Q1 as u8 => Ok(XXX::Q1),
                    _ if x == XXX::R1 as u8 => Ok(XXX::R1),
                    _ if x == XXX::S1 as u8 => Ok(XXX::S1),
                    _ if x == XXX::T1 as u8 => Ok(XXX::T1),
                    _ if x == XXX::U1 as u8 => Ok(XXX::U1),
                    _ if x == XXX::V1 as u8 => Ok(XXX::V1),
                    _ if x == XXX::W1 as u8 => Ok(XXX::W1),
                    _ if x == XXX::X1 as u8 => Ok(XXX::X1),
                    _ if x == XXX::Y1 as u8 => Ok(XXX::Y1),
                    _ if x == XXX::Z1 as u8 => Ok(XXX::Z1),
                    _ if x == XXX::A2 as u8 => Ok(XXX::A2),
                    _ if x == XXX::B2 as u8 => Ok(XXX::B2),
                    _ if x == XXX::C2 as u8 => Ok(XXX::C2),
                    _ if x == XXX::D2 as u8 => Ok(XXX::D2),
                    _ if x == XXX::E2 as u8 => Ok(XXX::E2),
                    _ if x == XXX::F2 as u8 => Ok(XXX::F2),
                    _ if x == XXX::G2 as u8 => Ok(XXX::G2),
                    _ if x == XXX::H2 as u8 => Ok(XXX::H2),
                    _ if x == XXX::I2 as u8 => Ok(XXX::I2),
                    _ if x == XXX::J2 as u8 => Ok(XXX::J2),
                    _ if x == XXX::K2 as u8 => Ok(XXX::K2),
                    _ if x == XXX::L2 as u8 => Ok(XXX::L2),
                    _ if x == XXX::M2 as u8 => Ok(XXX::M2),
                    _ if x == XXX::N2 as u8 => Ok(XXX::N2),
                    _ if x == XXX::O2 as u8 => Ok(XXX::O2),
                    _ if x == XXX::P2 as u8 => Ok(XXX::P2),
                    _ if x == XXX::Q2 as u8 => Ok(XXX::Q2),
                    _ if x == XXX::R2 as u8 => Ok(XXX::R2),
                    _ if x == XXX::S2 as u8 => Ok(XXX::S2),
                    _ if x == XXX::T2 as u8 => Ok(XXX::T2),
                    _ if x == XXX::U2 as u8 => Ok(XXX::U2),
                    _ if x == XXX::V2 as u8 => Ok(XXX::V2),
                    _ if x == XXX::W2 as u8 => Ok(XXX::W2),
                    _ if x == XXX::X2 as u8 => Ok(XXX::X2),
                    _ if x == XXX::Y2 as u8 => Ok(XXX::Y2),
                    _ if x == XXX::Z2 as u8 => Ok(XXX::Z2),
                    _ => Err(()),
                }
        }
}

Rust does a great job of boiling it down:

<XXX as core::convert::TryFrom<u8>>::try_from:
        cmp dil, 52
        mov eax, 52
        cmovb eax, edi
        ret

But it doesn't inline this properly when used in other functions:

impl XXX {
    pub fn new(x: u8) -> Self {
        Self::try_from(x).expect("bad num")
    }
}

generates:

XXX::new:
        push rax
        call qword ptr [rip + <XXX as core::convert::TryFrom<u8>>::try_from@GOTPCREL]
        cmp al, 52
        je .LBB139_1
        pop rcx
        ret

An explicit #[inline] given to try_from results in this instead:

XXX::new:
        cmp dil, 52
        jae .LBB139_1
        mov eax, edi
        ret

rustc --version --verbose:

rustc 1.74.1 (a28077b28 2023-12-04)
binary: rustc
commit-hash: a28077b28a02b92985b3a3faecf92813155f1ea1
commit-date: 2023-12-04
host: x86_64-unknown-linux-gnu
release: 1.74.1
LLVM version: 17.0.4

[saethlin]

Bumping up -Zcross-crate-inline-threshold=200 causes this to be inlined. Which is very strange to me (as the author of that flag) because this function is still not inlined in MIR and I don't see a codegen unit boundary between the function and call.

---

This line of thinking is irrelevant, cross-crate-inlining makes the function hidden, which makes it fastcc which is probably just a bonus to inlining. Not the root problem, just a way to paper over this specific example.

[saethlin]

Oh. Strange. The SimplifyCFG pass that would shrink this down to an instruction seems to be run after inlining??

[nikic]

The problem is likely that the transform doing this is the same one that produces switch lookup tables, which needs to happen post-inline. We could allow conversion of switch in cases where it does not produce a lookup table early.

[hanna-kruppe]

We could allow conversion of switch in cases where it does not produce a lookup table early.

Do you have an idea how difficult this would be? Also, is there already an issue on the LLVM side for it?

I'm currently running into confusing and sometimes bad codegen related to this and I don't think I'm doing something particularly unusual. Fieldless enums corresponding to small integers are useful for many purposes, and a big match that can be SimplifyCFG'd away is the safe and easy way to do the integer -> enum conversion without introducing tricky unsafe code. Quite a few popular libraries have macros/derives that generate such code (num_enum, strum, enum-ordinalize, enum-map, and probably more). If nothing else, there should be a nice compile time win for everything that heavily uses this pattern by cleaning it up before inlining.

In the code I'm currently working on, I'm also seeing missed optimizations that appear related to not doing this simplification ASAP. For example, mapping one enum with discriminants 0..3 to the values 7..10 of a larger enum via conversions to and from uN types ends up like this (Alive2 agrees it could just return add %self, 7):

define noundef range(i8 7, 10) i8 @small_enum_to_large_enum(i8 noundef range(i8 0, 3) %self) unnamed_addr #0 {
start:
  %narrow = add nuw nsw i8 %self, 7
  %trunc.i = trunc nuw i8 %narrow to i4
  switch i4 %trunc.i, label %default.unreachable [
    i4 -7, label %bb11.i
    i4 -8, label %bb12.i
    i4 7, label %large_enum_from_int.exit
  ]

default.unreachable:
  unreachable

bb12.i:
  br label %large_enum_from_int.exit

bb11.i:
  br label %large_enum_from_int.exit

large_enum_from_int.exit:
  %_0.sroa.0.0.i = phi i8 [ 8, %bb12.i ], [ 9, %bb11.i ], [ 7, %start ]
  ret i8 %_0.sroa.0.0.i
}

Obviously it would be best if the LUT construction in SimplifyCFG could handle the reduced switch properly (I guess that's llvm/llvm-project#64231 or something related). But currently it doesn't, and running SwitchToLookupTable before the case values get narrowed could help. In this specific case, it seems it wouldn't have helped anyway because of pass order (SROA is needed to clean up rustc's initial codegen, and then instcombine runs before SimplifyCFG gets another chance). But with more improvements along the lines of #138759 the first run of SimplifyCFG should already be able to turn the switch into an add.

Edit: actually, before IPO passes the switch has a reachable default case and can only turn into select (icmp ...), ... which is apparently also problematic to introduce this early, see llvm/llvm-project#53853