Extend switch condition in optimizeSwitchPhiConst when free

In a case like:

    switch((i32)x) { case 42: phi((i64)42, ...); }

replace `(i64)42` with `zext(x)` when we can do so for free.

This fixes a part of https://github.com/llvm/llvm-project/issues/55153

Differential Revision: https://reviews.llvm.org/D124897

1 files changed, 29 insertions, 4 deletions

diff --git a/llvm/lib/CodeGen/CodeGenPrepare.cpp b/llvm/lib/CodeGen/CodeGenPrepare.cpp
index d100038..9ed4f5b 100644
--- a/llvm/lib/CodeGen/CodeGenPrepare.cpp
+++ b/llvm/lib/CodeGen/CodeGenPrepare.cpp
@@ -7048,12 +7048,29 @@ bool CodeGenPrepare::optimizeSwitchPhiConstants(SwitchInst *SI) {
     bool CheckedForSinglePred = false;
     for (PHINode &PHI : CaseBB->phis()) {
       Type *PHIType = PHI.getType();
-      if (PHIType == ConditionType) {
+      // If ZExt is free then we can also catch patterns like this:
+      //   switch((i32)x) { case 42: phi((i64)42, ...); }
+      // and replace `(i64)42` with `zext i32 %x to i64`.
+      bool TryZExt =
+          PHIType->isIntegerTy() &&
+          PHIType->getIntegerBitWidth() > ConditionType->getIntegerBitWidth() &&
+          TLI->isZExtFree(ConditionType, PHIType);
+      if (PHIType == ConditionType || TryZExt) {
         // Set to true to skip this case because of multiple preds.
         bool SkipCase = false;
+        Value *Replacement = nullptr;
         for (unsigned I = 0, E = PHI.getNumIncomingValues(); I != E; I++) {
-          if (PHI.getIncomingValue(I) != CaseValue ||
-              PHI.getIncomingBlock(I) != SwitchBB)
+          Value *PHIValue = PHI.getIncomingValue(I);
+          if (PHIValue != CaseValue) {
+            if (!TryZExt)
+              continue;
+            ConstantInt *PHIValueInt = dyn_cast<ConstantInt>(PHIValue);
+            if (!PHIValueInt ||
+                PHIValueInt->getValue() !=
+                    CaseValue->getValue().zext(PHIType->getIntegerBitWidth()))
+              continue;
+          }
+          if (PHI.getIncomingBlock(I) != SwitchBB)
             continue;
           // We cannot optimize if there are multiple case labels jumping to
           // this block.  This check may get expensive when there are many
@@ -7066,7 +7083,15 @@ bool CodeGenPrepare::optimizeSwitchPhiConstants(SwitchInst *SI) {
             }
           }
 
-          PHI.setIncomingValue(I, Condition);
+          if (Replacement == nullptr) {
+            if (PHIValue == CaseValue) {
+              Replacement = Condition;
+            } else {
+              IRBuilder<> Builder(SI);
+              Replacement = Builder.CreateZExt(Condition, PHIType);
+            }
+          }
+          PHI.setIncomingValue(I, Replacement);
           Changed = true;
         }
         if (SkipCase)