switch not recognized as equivalent to sext

[scottmcm]

Opt transformation showing suboptimal current result on trunk: https://llvm.godbolt.org/z/jj6vcv5q3
Alive2 demonstration that the proposed transformation is correct: https://alive2.llvm.org/ce/z/hAAoh_
Original motivating example in rust-lang/rust#95313: https://rust.godbolt.org/z/dvYe9r78c

The following code:

define i32 @src(i8 noundef %0) {
start:
  %1 = alloca i32, align 4
  %order = alloca i8, align 1
  store i8 %0, i8* %order, align 1
  %_2 = load i8, i8* %order, align 1, !range !2, !noundef !3
  switch i8 %_2, label %bb2 [
    i8 -1, label %bb3
    i8 0, label %bb4
    i8 1, label %bb1
  ]

bb2:                                              ; preds = %start
  unreachable

bb3:                                              ; preds = %start
  store i32 -1, i32* %1, align 4
  br label %bb5

bb4:                                              ; preds = %start
  store i32 0, i32* %1, align 4
  br label %bb5

bb1:                                              ; preds = %start
  store i32 1, i32* %1, align 4
  br label %bb5

bb5:                                              ; preds = %bb3, %bb4, %bb1
  %2 = load i32, i32* %1, align 4
  ret i32 %2
}

!2 = !{i8 -1, i8 2}
!3 = !{}

Optimizes to this slightly suboptimal IR:

define i32 @src(i8 noundef %0) local_unnamed_addr #0 {
  %switch.tableidx = add i8 %0, 1
  %switch.idx.cast = zext i8 %switch.tableidx to i32
  %switch.offset = add nsw i32 %switch.idx.cast, -1
  ret i32 %switch.offset
}

But as Alive confirms, it could just be a sext:

define i32 @tgt(i8 noundef %0) {
%start:
  %1 = sext i8 noundef %0 to i32
  ret i32 %1
}
// Transformation seems to be correct!

[scottmcm]

I had to use different types because of integer promotion in C++, but here's a clang repro too:

long long square(int num) {
    switch (num) {
        case -1: return -1;
        case 0: return 0;
        case +1: return +1;
        default: __builtin_unreachable();
    }
}

https://clang.godbolt.org/z/57jqWxbxe

square(int):                             # @square(int)
        lea     eax, [rdi + 1]
        dec     rax
        ret

[nikic]

I think we can extend

llvm-project/llvm/lib/Transforms/InstCombine/InstCombinePHI.cpp

Line 1264 in 6c2be88

static Value *simplifyUsingControlFlow(InstCombiner &Self, PHINode &PN,

to handle this case. It already deals with an invert transform, could easily support sext and zext as well.

[hkmatsumoto]

IMHO only extending simplifyUsingControlFlow does not work because the LLVM IR in question does not contain PHI nodes/functions.

Is it a nice approach to transform the IR into something like this containing PHI node, and let existing optimization passes to make further optimizations?

[nikic]

@hkmatsumoto SROA will convert the load/stores into a phi node.

[hkmatsumoto]

Oh, I'm convinced. Let me see if I can make a patch for this.

[hkmatsumoto]

Opened https://reviews.llvm.org/D122913