Extend switch condition in optimizeSwitchPhiConst when free

[MatzeB]

Summary

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

[spatel]

Push the previous patches in this sequence, so we can get some bot and real-world testing underway before adding complexity to the transform?

[spatel]

The previous patches seem to be stable now, so LGTM with a couple of minor improvements.

llvm/lib/CodeGen/CodeGenPrepare.cpp

7027

formatting: TryZext or TryZExt

7035–7038

Prefer dyn_cast to isa + cast, so something like this:

auto *PHIValue = dyn_cast<ConstantInt>(PHI.getIncomingValue(I));
if (PHIValue != CaseValue)
  if (!tryZExt || !PHIValue ||
      PHIValue->getValue() !=
          CaseValue->getValue().zext(PHIType->getIntegerBitWidth()))
    continue;

llvm/lib/CodeGen/CodeGenPrepare.cpp

bool CodeGenPrepare::optimizeSwitchPhiConstants(SwitchInst *SI) {
   Type *ConditionType = Condition->getType();
 
   for (const SwitchInst::CaseHandle &Case : SI->cases()) {
     bool CheckedCase = false;
     ConstantInt *CaseValue = Case.getCaseValue();
     BasicBlock *CaseBB = Case.getCaseSuccessor();
     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) {
         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) {
-            // We cannot optimize if there are multiple case labels jumping to
-            // this block.  This check may get expensive when there are many
-            // case labels so we test for it last.
-            if (!CheckedCase) {
-              CheckedCase = true;
-              if (SI->findCaseDest(CaseBB) == nullptr) {
-                SkipCase = true;
-                break;
-              }
-            }
+          Value *PHIValue = PHI.getIncomingValue(I);
+          if (PHIValue != CaseValue &&
+              (!tryZExt || !isa<ConstantInt>(PHIValue) ||
+               cast<ConstantInt>(PHIValue)->getValue() !=
+                   CaseValue->getValue().zext(PHIType->getIntegerBitWidth())))
+            continue;
+          if (PHI.getIncomingBlock(I) != SwitchBB)
+            continue;
 
-            PHI.setIncomingValue(I, Condition);
-            Changed = true;
+          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)
           break;
       }
     }
   }
   return Changed;
 }

llvm/test/CodeGen/X86/switch-phi-const.ll

 }
 
 @g64 = global i64 0
 @effect64 = global i64 0
 
 define void @switch_trunc_phi_const(i32 %x) {
 ; CHECK-LABEL: switch_trunc_phi_const:
 ; CHECK:       # %bb.0: # %bb0
-; CHECK-NEXT:    movzbl %dil, %r8d
-; CHECK-NEXT:    movl $3895, %ecx # imm = 0xF37
-; CHECK-NEXT:    movl $42, %esi
-; CHECK-NEXT:    movl $13, %edx
-; CHECK-NEXT:    movl $5, %edi
-; CHECK-NEXT:    movl $1, %eax
-; CHECK-NEXT:    decl %r8d
-; CHECK-NEXT:    cmpl $54, %r8d
+; CHECK-NEXT:    # kill: def $edi killed $edi def $rdi
+; CHECK-NEXT:    movzbl %dil, %ecx
+; CHECK-NEXT:    movzbl %dil, %eax
+; CHECK-NEXT:    movl $3895, %edx # imm = 0xF37
+; CHECK-NEXT:    decl %ecx
+; CHECK-NEXT:    cmpl $54, %ecx
 ; CHECK-NEXT:    ja .LBB1_8
 ; CHECK-NEXT:  # %bb.1: # %bb0
-; CHECK-NEXT:    jmpq *.LJTI1_0(,%r8,8)
+; CHECK-NEXT:    jmpq *.LJTI1_0(,%rcx,8)
 ; CHECK-NEXT:  .LBB1_8: # %default
 ; CHECK-NEXT:    retq
 ; CHECK-NEXT:  .LBB1_2: # %case_1_loop
 ; CHECK-NEXT:    movq effect64@GOTPCREL(%rip), %rcx
 ; CHECK-NEXT:    movq $1, (%rcx)
-; CHECK-NEXT:    movq %rax, %rdi
 ; CHECK-NEXT:  .LBB1_3: # %case_5
-; CHECK-NEXT:    movq effect64@GOTPCREL(%rip), %rax
-; CHECK-NEXT:    movq $5, (%rax)
-; CHECK-NEXT:    movq %rdi, %rdx
+; CHECK-NEXT:    movq effect64@GOTPCREL(%rip), %rcx
+; CHECK-NEXT:    movq $5, (%rcx)
 ; CHECK-NEXT:  .LBB1_4: # %case_13
-; CHECK-NEXT:    movq effect64@GOTPCREL(%rip), %rax
-; CHECK-NEXT:    movq $13, (%rax)
-; CHECK-NEXT:    movq %rdx, %rsi
+; CHECK-NEXT:    movq effect64@GOTPCREL(%rip), %rcx
+; CHECK-NEXT:    movq $13, (%rcx)
 ; CHECK-NEXT:  .LBB1_5: # %case_42
-; CHECK-NEXT:    movq effect64@GOTPCREL(%rip), %rax
-; CHECK-NEXT:    movq %rsi, (%rax)
-; CHECK-NEXT:    movl $55, %ecx
+; CHECK-NEXT:    movq effect64@GOTPCREL(%rip), %rcx
+; CHECK-NEXT:    movq %rax, (%rcx)
+; CHECK-NEXT:    movl $55, %edx
 ; CHECK-NEXT:  .LBB1_6: # %case_55
 ; CHECK-NEXT:    movq effect64@GOTPCREL(%rip), %rax
-; CHECK-NEXT:    movq %rcx, (%rax)
+; CHECK-NEXT:    movq %rdx, (%rax)
 ; CHECK-NEXT:  .LBB1_7: # %case_7
 ; CHECK-NEXT:    movq g64@GOTPCREL(%rip), %rax
 ; CHECK-NEXT:    movq (%rax), %rax
 ; CHECK-NEXT:    movq effect64@GOTPCREL(%rip), %rcx
 ; CHECK-NEXT:    movq $7, (%rcx)
 ; CHECK-NEXT:    jmp .LBB1_2
 bb0:
   %x_trunc = trunc i32 %x to i8

llvm/test/Transforms/CodeGenPrepare/X86/switch-phi-const.ll

 case_7:
   %x_loopback = load i32, i32* @g, align 4
   store i32 7, i32* @effect, align 4
   br label %case_13
 
 default:
   ret void
 }
+
+@g64 = global i64 0
+@effect64 = global i64 0
+
+define void @switch_trunc_phi_const(i32 %x) {
+; CHECK-LABEL: @switch_trunc_phi_const(
+; CHECK-NEXT:  bb0:
+; CHECK-NEXT:    [[TMP0:%.*]] = zext i32 [[X:%.*]] to i64
+; CHECK-NEXT:    [[TMP1:%.*]] = zext i32 [[X]] to i64
+; CHECK-NEXT:    switch i32 [[X]], label [[DEFAULT:%.*]] [
+; CHECK-NEXT:    i32 13, label [[CASE_13:%.*]]
+; CHECK-NEXT:    i32 42, label [[CASE_42:%.*]]
+; CHECK-NEXT:    i32 55, label [[CASE_55:%.*]]
+; CHECK-NEXT:    i32 7, label [[CASE_7:%.*]]
+; CHECK-NEXT:    ]
+; CHECK:       case_13:
+; CHECK-NEXT:    [[X0:%.*]] = phi i64 [ [[TMP0]], [[BB0:%.*]] ], [ [[X7:%.*]], [[CASE_7]] ]
+; CHECK-NEXT:    store i64 13, i64* @effect64, align 4
+; CHECK-NEXT:    br label [[CASE_42]]
+; CHECK:       case_42:
+; CHECK-NEXT:    [[X1:%.*]] = phi i64 [ [[TMP1]], [[BB0]] ], [ [[X0]], [[CASE_13]] ]
+; CHECK-NEXT:    store i64 [[X1]], i64* @effect64, align 4
+; CHECK-NEXT:    br label [[CASE_55]]
+; CHECK:       case_55:
+; CHECK-NEXT:    [[X2:%.*]] = phi i64 [ 3895, [[BB0]] ], [ 55, [[CASE_42]] ]
+; CHECK-NEXT:    store i64 [[X2]], i64* @effect64, align 4
+; CHECK-NEXT:    br label [[DEFAULT]]
+; CHECK:       case_7:
+; CHECK-NEXT:    [[X7]] = load i64, i64* @g64, align 4
+; CHECK-NEXT:    store i64 7, i64* @effect64, align 4
+; CHECK-NEXT:    br label [[CASE_13]]
+; CHECK:       default:
+; CHECK-NEXT:    ret void
+;
+bb0:
+  switch i32 %x, label %default [
+  i32 13, label %case_13
+  i32 42, label %case_42
+  i32 55, label %case_55
+  i32 7, label %case_7
+  ]
+
+case_13:
+  ; We should replace 13 with zext %x to i64
+  %x0 = phi i64 [ 13, %bb0 ], [ %x7, %case_7 ]
+  store i64 13, i64* @effect64, align 4
+  br label %case_42
+
+case_42:
+  ; We should replace 42 with zext i32 %x to i64
+  %x1 = phi i64 [ 42, %bb0 ], [ %x0, %case_13 ]
+  store i64 %x1, i64* @effect64, align 4
+  br label %case_55
+
+case_55:
+  ; We must not replace any of the PHI arguments! (3898 == 0xf00 + 55)
+  %x2 = phi i64 [ 3895, %bb0 ], [ 55, %case_42 ]
+  store i64 %x2, i64* @effect64, align 4
+  br label %default
+
+case_7:
+  %x7 = load i64, i64* @g64, align 4
+  store i64 7, i64* @effect64, align 4
+  br label %case_13
+
+default:
+  ret void
+}