[InstCombine] Fold `switch(zext/sext(X))` into `switch(X)` (#76988)

This patch folds `switch(zext/sext(X))` into `switch(X)`.
The original motivation of this patch is to optimize a pattern found in
cvc5. For example:
```
  %bf.load.i = load i16, ptr %d_kind.i, align 8
  %bf.clear.i = and i16 %bf.load.i, 1023
  %bf.cast.i = zext nneg i16 %bf.clear.i to i32
  switch i32 %bf.cast.i, label %if.else [
    i32 335, label %if.then
    i32 303, label %if.then
  ]

if.then:                                          ; preds = %entry, %entry
  %d_children.i.i = getelementptr inbounds %"class.cvc5::internal::expr::NodeValue", ptr %0, i64 0, i32 3
  %cmp.i.i.i.i.i = icmp eq i16 %bf.clear.i, 1023
  %cond.i.i.i.i.i = select i1 %cmp.i.i.i.i.i, i32 -1, i32 %bf.cast.i
```
`%cmp.i.i.i.i.i` always evaluates to false because `%bf.clear.i` can
only be 335 or 303.
Folding `switch i32 %bf.cast.i` to `switch i16 %bf.clear.i` will help
`CVP` to handle this case.
See also
https://github.com/llvm/llvm-project/pull/76928#issuecomment-1877055722.

Compile-time impact:
http://llvm-compile-time-tracker.com/compare.php?from=7954c57124b495fbdc73674d71f2e366e4afe522&to=502b13ed34e561d995ae1f724cf06d20008bd86f&stat=instructions:u

|stage1-O3|stage1-ReleaseThinLTO|stage1-ReleaseLTO-g|stage1-O0-g|stage2-O3|stage2-O0-g|stage2-clang|
|--|--|--|--|--|--|--|
|+0.03%|+0.06%|+0.07%|+0.00%|-0.02%|-0.03%|+0.02%|

4 files changed, 202 insertions, 17 deletions

diff --git a/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp b/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
index e845bf38df2d..7f2018b3a199 100644
--- a/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
@@ -3247,6 +3247,25 @@ Instruction *InstCombinerImpl::visitSwitchInst(SwitchInst &SI) {
     }
   }
 
+  // Fold switch(zext/sext(X)) into switch(X) if possible.
+  if (match(Cond, m_ZExtOrSExt(m_Value(Op0)))) {
+    bool IsZExt = isa<ZExtInst>(Cond);
+    Type *SrcTy = Op0->getType();
+    unsigned NewWidth = SrcTy->getScalarSizeInBits();
+
+    if (all_of(SI.cases(), [&](const auto &Case) {
+          const APInt &CaseVal = Case.getCaseValue()->getValue();
+          return IsZExt ? CaseVal.isIntN(NewWidth)
+                        : CaseVal.isSignedIntN(NewWidth);
+        })) {
+      for (auto &Case : SI.cases()) {
+        APInt TruncatedCase = Case.getCaseValue()->getValue().trunc(NewWidth);
+        Case.setValue(ConstantInt::get(SI.getContext(), TruncatedCase));
+      }
+      return replaceOperand(SI, 0, Op0);
+    }
+  }
+
   KnownBits Known = computeKnownBits(Cond, 0, &SI);
   unsigned LeadingKnownZeros = Known.countMinLeadingZeros();
   unsigned LeadingKnownOnes = Known.countMinLeadingOnes();
diff --git a/llvm/test/Transforms/InstCombine/phi.ll b/llvm/test/Transforms/InstCombine/phi.ll
index 90818771675b..717f4c682a15 100644
--- a/llvm/test/Transforms/InstCombine/phi.ll
+++ b/llvm/test/Transforms/InstCombine/phi.ll
@@ -996,10 +996,9 @@ done:
 define i1 @PR24766(i8 %x1, i8 %x2, i8 %condition) {
 ; CHECK-LABEL: @PR24766(
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[CONV:%.*]] = sext i8 [[CONDITION:%.*]] to i32
-; CHECK-NEXT:    switch i32 [[CONV]], label [[EPILOG:%.*]] [
-; CHECK-NEXT:      i32 0, label [[SW1:%.*]]
-; CHECK-NEXT:      i32 1, label [[SW2:%.*]]
+; CHECK-NEXT:    switch i8 [[CONDITION:%.*]], label [[EPILOG:%.*]] [
+; CHECK-NEXT:      i8 0, label [[SW1:%.*]]
+; CHECK-NEXT:      i8 1, label [[SW2:%.*]]
 ; CHECK-NEXT:    ]
 ; CHECK:       sw1:
 ; CHECK-NEXT:    [[CMP1:%.*]] = icmp eq i8 [[X1:%.*]], [[X2:%.*]]
@@ -1040,10 +1039,9 @@ epilog:
 define i1 @PR24766_no_constants(i8 %x1, i8 %x2, i8 %condition, i1 %another_condition) {
 ; CHECK-LABEL: @PR24766_no_constants(
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[CONV:%.*]] = sext i8 [[CONDITION:%.*]] to i32
-; CHECK-NEXT:    switch i32 [[CONV]], label [[EPILOG:%.*]] [
-; CHECK-NEXT:      i32 0, label [[SW1:%.*]]
-; CHECK-NEXT:      i32 1, label [[SW2:%.*]]
+; CHECK-NEXT:    switch i8 [[CONDITION:%.*]], label [[EPILOG:%.*]] [
+; CHECK-NEXT:      i8 0, label [[SW1:%.*]]
+; CHECK-NEXT:      i8 1, label [[SW2:%.*]]
 ; CHECK-NEXT:    ]
 ; CHECK:       sw1:
 ; CHECK-NEXT:    [[CMP1:%.*]] = icmp eq i8 [[X1:%.*]], [[X2:%.*]]
@@ -1085,10 +1083,9 @@ epilog:
 define i1 @PR24766_two_constants(i8 %x1, i8 %x2, i8 %condition) {
 ; CHECK-LABEL: @PR24766_two_constants(
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[CONV:%.*]] = sext i8 [[CONDITION:%.*]] to i32
-; CHECK-NEXT:    switch i32 [[CONV]], label [[EPILOG:%.*]] [
-; CHECK-NEXT:      i32 0, label [[SW1:%.*]]
-; CHECK-NEXT:      i32 1, label [[SW2:%.*]]
+; CHECK-NEXT:    switch i8 [[CONDITION:%.*]], label [[EPILOG:%.*]] [
+; CHECK-NEXT:      i8 0, label [[SW1:%.*]]
+; CHECK-NEXT:      i8 1, label [[SW2:%.*]]
 ; CHECK-NEXT:    ]
 ; CHECK:       sw1:
 ; CHECK-NEXT:    [[CMP1:%.*]] = icmp eq i8 [[X1:%.*]], [[X2:%.*]]
@@ -1128,11 +1125,10 @@ epilog:
 define i1 @PR24766_two_constants_two_var(i8 %x1, i8 %x2, i8 %condition) {
 ; CHECK-LABEL: @PR24766_two_constants_two_var(
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[CONV:%.*]] = sext i8 [[CONDITION:%.*]] to i32
-; CHECK-NEXT:    switch i32 [[CONV]], label [[EPILOG:%.*]] [
-; CHECK-NEXT:      i32 0, label [[SW1:%.*]]
-; CHECK-NEXT:      i32 1, label [[SW2:%.*]]
-; CHECK-NEXT:      i32 2, label [[SW3:%.*]]
+; CHECK-NEXT:    switch i8 [[CONDITION:%.*]], label [[EPILOG:%.*]] [
+; CHECK-NEXT:      i8 0, label [[SW1:%.*]]
+; CHECK-NEXT:      i8 1, label [[SW2:%.*]]
+; CHECK-NEXT:      i8 2, label [[SW3:%.*]]
 ; CHECK-NEXT:    ]
 ; CHECK:       sw1:
 ; CHECK-NEXT:    [[CMP1:%.*]] = icmp eq i8 [[X1:%.*]], [[X2:%.*]]
diff --git a/llvm/test/Transforms/InstCombine/switch-zext-sext.ll b/llvm/test/Transforms/InstCombine/switch-zext-sext.ll
new file mode 100644
index 000000000000..c09441352acc
--- /dev/null
+++ b/llvm/test/Transforms/InstCombine/switch-zext-sext.ll
@@ -0,0 +1,122 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 4
+; RUN: opt %s -passes=instcombine -S | FileCheck %s
+
+define i1 @test_switch_with_zext(i16 %a, i1 %b, i1 %c) {
+; CHECK-LABEL: define i1 @test_switch_with_zext(
+; CHECK-SAME: i16 [[A:%.*]], i1 [[B:%.*]], i1 [[C:%.*]]) {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    switch i16 [[A]], label [[SW_DEFAULT:%.*]] [
+; CHECK-NEXT:      i16 37, label [[SW_BB:%.*]]
+; CHECK-NEXT:      i16 38, label [[SW_BB]]
+; CHECK-NEXT:      i16 39, label [[SW_BB]]
+; CHECK-NEXT:    ]
+; CHECK:       sw.bb:
+; CHECK-NEXT:    ret i1 [[B]]
+; CHECK:       sw.default:
+; CHECK-NEXT:    ret i1 [[C]]
+;
+entry:
+  %a.ext = zext i16 %a to i32
+  switch i32 %a.ext, label %sw.default [
+  i32 37, label %sw.bb
+  i32 38, label %sw.bb
+  i32 39, label %sw.bb
+  ]
+
+sw.bb:
+  ret i1 %b
+sw.default:
+  ret i1 %c
+}
+
+define i1 @test_switch_with_sext(i16 %a, i1 %b, i1 %c) {
+; CHECK-LABEL: define i1 @test_switch_with_sext(
+; CHECK-SAME: i16 [[A:%.*]], i1 [[B:%.*]], i1 [[C:%.*]]) {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    switch i16 [[A]], label [[SW_DEFAULT:%.*]] [
+; CHECK-NEXT:      i16 37, label [[SW_BB:%.*]]
+; CHECK-NEXT:      i16 38, label [[SW_BB]]
+; CHECK-NEXT:      i16 39, label [[SW_BB]]
+; CHECK-NEXT:    ]
+; CHECK:       sw.bb:
+; CHECK-NEXT:    ret i1 [[B]]
+; CHECK:       sw.default:
+; CHECK-NEXT:    ret i1 [[C]]
+;
+entry:
+  %a.ext = sext i16 %a to i32
+  switch i32 %a.ext, label %sw.default [
+  i32 37, label %sw.bb
+  i32 38, label %sw.bb
+  i32 39, label %sw.bb
+  ]
+
+sw.bb:
+  ret i1 %b
+sw.default:
+  ret i1 %c
+}
+
+; Negative tests
+
+define i1 @test_switch_with_zext_unreachable_case(i16 %a, i1 %b, i1 %c) {
+; CHECK-LABEL: define i1 @test_switch_with_zext_unreachable_case(
+; CHECK-SAME: i16 [[A:%.*]], i1 [[B:%.*]], i1 [[C:%.*]]) {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[A_EXT:%.*]] = zext i16 [[A]] to i32
+; CHECK-NEXT:    switch i32 [[A_EXT]], label [[SW_DEFAULT:%.*]] [
+; CHECK-NEXT:      i32 37, label [[SW_BB:%.*]]
+; CHECK-NEXT:      i32 38, label [[SW_BB]]
+; CHECK-NEXT:      i32 39, label [[SW_BB]]
+; CHECK-NEXT:      i32 65537, label [[SW_BB]]
+; CHECK-NEXT:    ]
+; CHECK:       sw.bb:
+; CHECK-NEXT:    ret i1 [[B]]
+; CHECK:       sw.default:
+; CHECK-NEXT:    ret i1 [[C]]
+;
+entry:
+  %a.ext = zext i16 %a to i32
+  switch i32 %a.ext, label %sw.default [
+  i32 37, label %sw.bb
+  i32 38, label %sw.bb
+  i32 39, label %sw.bb
+  i32 65537, label %sw.bb
+  ]
+
+sw.bb:
+  ret i1 %b
+sw.default:
+  ret i1 %c
+}
+
+define i1 @test_switch_with_sext_unreachable_case(i16 %a, i1 %b, i1 %c) {
+; CHECK-LABEL: define i1 @test_switch_with_sext_unreachable_case(
+; CHECK-SAME: i16 [[A:%.*]], i1 [[B:%.*]], i1 [[C:%.*]]) {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[A_EXT:%.*]] = sext i16 [[A]] to i32
+; CHECK-NEXT:    switch i32 [[A_EXT]], label [[SW_DEFAULT:%.*]] [
+; CHECK-NEXT:      i32 37, label [[SW_BB:%.*]]
+; CHECK-NEXT:      i32 38, label [[SW_BB]]
+; CHECK-NEXT:      i32 39, label [[SW_BB]]
+; CHECK-NEXT:      i32 -65537, label [[SW_BB]]
+; CHECK-NEXT:    ]
+; CHECK:       sw.bb:
+; CHECK-NEXT:    ret i1 [[B]]
+; CHECK:       sw.default:
+; CHECK-NEXT:    ret i1 [[C]]
+;
+entry:
+  %a.ext = sext i16 %a to i32
+  switch i32 %a.ext, label %sw.default [
+  i32 37, label %sw.bb
+  i32 38, label %sw.bb
+  i32 39, label %sw.bb
+  i32 -65537, label %sw.bb
+  ]
+
+sw.bb:
+  ret i1 %b
+sw.default:
+  ret i1 %c
+}
diff --git a/llvm/test/Transforms/PhaseOrdering/switch-sext.ll b/llvm/test/Transforms/PhaseOrdering/switch-sext.ll
new file mode 100644
index 000000000000..e39771b629e0
--- /dev/null
+++ b/llvm/test/Transforms/PhaseOrdering/switch-sext.ll
@@ -0,0 +1,48 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 4
+; RUN: opt -S -passes='default<O3>' < %s | FileCheck %s
+
+define i8 @test_switch_with_sext_phi(i8 %code) {
+; CHECK-LABEL: define noundef i8 @test_switch_with_sext_phi(
+; CHECK-SAME: i8 [[CODE:%.*]]) local_unnamed_addr #[[ATTR0:[0-9]+]] {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    switch i8 [[CODE]], label [[SW_EPILOG:%.*]] [
+; CHECK-NEXT:      i8 76, label [[SW_BB3:%.*]]
+; CHECK-NEXT:      i8 108, label [[SW_BB2:%.*]]
+; CHECK-NEXT:    ]
+; CHECK:       sw.bb2:
+; CHECK-NEXT:    br label [[SW_EPILOG]]
+; CHECK:       sw.bb3:
+; CHECK-NEXT:    br label [[SW_EPILOG]]
+; CHECK:       sw.epilog:
+; CHECK-NEXT:    [[PEP_CODE:%.*]] = phi i8 [ 81, [[SW_BB3]] ], [ 113, [[SW_BB2]] ], [ [[CODE]], [[ENTRY:%.*]] ]
+; CHECK-NEXT:    ret i8 [[PEP_CODE]]
+;
+entry:
+  %conv = sext i8 %code to i32
+  switch i32 %conv, label %sw.default [
+  i32 105, label %sw.epilog
+  i32 73, label %sw.bb1
+  i32 108, label %sw.bb2
+  i32 76, label %sw.bb3
+  i32 63, label %sw.bb4
+  ]
+
+sw.bb1:
+  br label %sw.epilog
+
+sw.bb2:
+  br label %sw.epilog
+
+sw.bb3:
+  br label %sw.epilog
+
+sw.bb4:
+  br label %sw.epilog
+
+sw.default:
+  br label %sw.epilog
+
+sw.epilog:
+  %pep_code = phi i8 [ %code, %sw.default ], [ 63, %sw.bb4 ], [ 81, %sw.bb3 ], [ 113, %sw.bb2 ], [ 73, %sw.bb1 ], [ 105, %entry ]
+  ret i8 %pep_code
+}