[CGP] widen switch condition and case constants to target's register width

[spatel]

Summary

This is a follow-up from the discussion in D12965. The block-at-a-time limitation of SelectionDAG also came up in D13297.

Without the InstCombine change from D12965, I don't expect this patch to make any difference in the real world because InstCombine will already be widening cases like this in visitSwitchInst(). But we need to have this CGP safety harness in place before proceeding with any shrinkage in D12965, so we won't generate extra extends for compares.

There are regression tests for CGP in both test/Transforms/CodeGenPrepare and test/CodeGen/* . I opted for IR regression tests in the patch because that seems like a clearer way to test the transform, but PowerPC CodeGen for the i16 widening test is shown below. x86 will need more work to solve: https://llvm.org/bugs/show_bug.cgi?id=22473

Before:

1. BB#0: mr 4, 3 extsh. 3, 4 ble 0, .LBB0_5
2. BB#1: cmpwi 3, 99 bgt 0, .LBB0_9
3. BB#2: rlwinm 4, 4, 0, 16, 31 <--- 32-bit mask/extend li 3, 0 cmplwi 4, 1 beqlr 0
4. BB#3: cmplwi 4, 10 bne 0, .LBB0_12
5. BB#4: li 3, 1 blr .LBB0_5: rlwinm 3, 4, 0, 16, 31 <--- 32-bit mask/extend cmplwi 3, 65436 beq 0, .LBB0_13
6. BB#6: cmplwi 3, 65526 beq 0, .LBB0_15
7. BB#7: cmplwi 3, 65535 bne 0, .LBB0_12
8. BB#8: li 3, 4 blr .LBB0_9: rlwinm 3, 4, 0, 16, 31 <--- 32-bit mask/extend cmplwi 3, 100 beq 0, .LBB0_14 ...

After:

1. 1. BB#0: rlwinm 4, 3, 0, 16, 31 <--- mask/extend to 32-bit and then use that for comparisons cmpwi 4, 999 ble 0, .LBB0_5
   2. BB#1: lis 3, 0 ori 3, 3, 65525 cmpw 4, 3 bgt 0, .LBB0_9
   3. BB#2: cmplwi 4, 1000 beq 0, .LBB0_14
2. BB#3: cmplwi 4, 65436 bne 0, .LBB0_13
   1. BB#4: li 3, 6 blr .LBB0_5: li 3, 0 cmplwi 4, 1 beqlr 0
   2. BB#6: cmplwi 4, 10 beq 0, .LBB0_12
   3. BB#7: cmplwi 4, 100 bne 0, .LBB0_13
   4. BB#8: li 3, 2 blr .LBB0_9: cmplwi 4, 65526 beq 0, .LBB0_15
   5. BB#10: cmplwi 4, 65535 bne 0, .LBB0_13 ...

[spatel]

Phab formatting was thrown off by '#' on the BB's. Removed below:

Before:

 BB#0:
  mr 4, 3
  extsh. 3, 4
  ble 0, .LBB0_5
 BB#1: 
  cmpwi	 3, 99
  bgt	 0, .LBB0_9
 BB#2:            
  rlwinm 4, 4, 0, 16, 31      <--- 32-bit mask/extend
  li 3, 0
  cmplwi	 4, 1
  beqlr 0
 BB#3:            
  cmplwi	 4, 10
  bne	 0, .LBB0_12
 BB#4:                      
  li 3, 1
  blr
.LBB0_5:                             
  rlwinm 3, 4, 0, 16, 31         <--- 32-bit mask/extend
  cmplwi	 3, 65436
  beq	 0, .LBB0_13
 BB#6:                            
  cmplwi	 3, 65526
  beq	 0, .LBB0_15
 BB#7:                       
  cmplwi	 3, 65535
  bne	 0, .LBB0_12
 BB#8:                       
  li 3, 4
  blr
.LBB0_9:                       
  rlwinm 3, 4, 0, 16, 31      <--- 32-bit mask/extend
  cmplwi	 3, 100
  beq	 0, .LBB0_14
...

After:

 BB#0:        
  rlwinm 4, 3, 0, 16, 31   <--- mask/extend to 32-bit and then use that for comparisons
  cmpwi	 4, 999
  ble 0, .LBB0_5
 BB#1:          
  lis 3, 0
  ori 3, 3, 65525
  cmpw	 4, 3
  bgt	 0, .LBB0_9
 BB#2:         
  cmplwi	 4, 1000
  beq	 0, .LBB0_14
 BB#3:    
  cmplwi	 4, 65436
  bne	 0, .LBB0_13
 BB#4:       
  li 3, 6
  blr
.LBB0_5:   
  li 3, 0
  cmplwi	 4, 1
  beqlr 0
 BB#6: 
  cmplwi	 4, 10
  beq	 0, .LBB0_12
 BB#7:             
  cmplwi	 4, 100
  bne	 0, .LBB0_13
 BB#8:             
  li 3, 2
  blr
.LBB0_9:       
  cmplwi	 4, 65526
  beq	 0, .LBB0_15
 BB#10:      
  cmplwi	 4, 65535
  bne	 0, .LBB0_13
...

[spatel]

Ping.

[reames]

This looks entirely reasonable to me, but I'm not an expert in this area. The code LGTM, but I'm not sure about the implications of the change. Can someone else comment on that?

lib/CodeGen/CodeGenPrepare.cpp

4015 ↗

(On Diff #36790)

auto

[spatel]

Ping * 2.

[hfinkel]

In D13532#262174, @spatel wrote:

After:

BB#0:        
 rlwinm 4, 3, 0, 16, 31   <--- mask/extend to 32-bit and then use that for comparisons
 cmpwi	 4, 999

...

Someone amusingly, even though you can't tell from this test case, this is somewhat suboptimal. Consider this:

$ cat /tmp/foo.c 
short foo(short a) { return a; }

$ clang -O3 -S -emit-llvm -o - /tmp/foo.c 
target datalayout = "E-m:e-i64:64-n32:64"
target triple = "powerpc64-unknown-linux-gnu"

define signext i16 @foo(i16 signext %a) #0 {
entry:
  ret i16 %a
}

and, notice here that the argument has the 'signext' attribute. This argument will be carried in an i32 register, but sign extended from i16. Thus, at least in theory, if we sign extended in this transformation, instead of zero-extended, we'd not actually need any extension instruction at all.

Can you think of any way to have it pick sext instead of zext when we know the input is really sign extended anyway? This might only apply to function arguments?

lib/CodeGen/CodeGenPrepare.cpp

4016 ↗

(On Diff #36790)

We don't need the 'OrBitCast' here, do we?

[spatel]

lib/CodeGen/CodeGenPrepare.cpp

4016 ↗

(On Diff #36790)

Correct - the check above guarantees that RegWidth is greater than the original type.

[spatel]

In D13532#276839, @hfinkel wrote:

Can you think of any way to have it pick sext instead of zext when we know the input is really sign extended anyway? This might only apply to function arguments?

Nice catch. I agree that we should sext instead of zext in that case. If it's only applicable to function args, it seems straightforward to add that check. New patch coming soon.

[spatel]

Patch updated:

1. Add check for a sign-extended function argument - in that case sext everything instead of zext.
2. Minimized tests - we don't need so many switch cases to show the differences in IR (the larger examples were better for visualizing the asm differences).
3. Added test case for #1.

[hfinkel]

LGTM, thanks!

llvm/trunk/lib/CodeGen/CodeGenPrepare.cpp

private:
     bool optimizeMemoryInst(Instruction *I, Value *Addr,
                             Type *AccessTy, unsigned AS);
     bool optimizeInlineAsmInst(CallInst *CS);
     bool optimizeCallInst(CallInst *CI, bool& ModifiedDT);
     bool moveExtToFormExtLoad(Instruction *&I);
     bool optimizeExtUses(Instruction *I);
     bool optimizeSelectInst(SelectInst *SI);
     bool optimizeShuffleVectorInst(ShuffleVectorInst *SI);
+    bool optimizeSwitchInst(SwitchInst *CI);
     bool optimizeExtractElementInst(Instruction *Inst);
     bool dupRetToEnableTailCallOpts(BasicBlock *BB);
     bool placeDbgValues(Function &F);
     bool sinkAndCmp(Function &F);
     bool extLdPromotion(TypePromotionTransaction &TPT, LoadInst *&LI,
                         Instruction *&Inst,
                         const SmallVectorImpl<Instruction *> &Exts,
                         unsigned CreatedInstCost);
bool CodeGenPrepare::optimizeShuffleVectorInst(ShuffleVectorInst *SVI) {
   if (SVI->use_empty()) {
     SVI->eraseFromParent();
     MadeChange = true;
   }
 
   return MadeChange;
 }
 
+bool CodeGenPrepare::optimizeSwitchInst(SwitchInst *SI) {
+  if (!TLI || !DL)
+    return false;
+
+  Value *Cond = SI->getCondition();
+  Type *OldType = Cond->getType();
+  LLVMContext &Context = Cond->getContext();
+  MVT RegType = TLI->getRegisterType(Context, TLI->getValueType(*DL, OldType));
+  unsigned RegWidth = RegType.getSizeInBits();
+
+  if (RegWidth <= cast<IntegerType>(OldType)->getBitWidth())
+    return false;
+
+  // If the register width is greater than the type width, expand the condition
+  // of the switch instruction and each case constant to the width of the
+  // register. By widening the type of the switch condition, subsequent
+  // comparisons (for case comparisons) will not need to be extended to the
+  // preferred register width, so we will potentially eliminate N-1 extends,
+  // where N is the number of cases in the switch.
+  auto *NewType = Type::getIntNTy(Context, RegWidth);
+
+  // Zero-extend the switch condition and case constants unless the switch
+  // condition is a function argument that is already being sign-extended.
+  // In that case, we can avoid an unnecessary mask/extension by sign-extending
+  // everything instead.
+  Instruction::CastOps ExtType = Instruction::ZExt;
+  if (auto *Arg = dyn_cast<Argument>(Cond))
+    if (Arg->hasSExtAttr())
+      ExtType = Instruction::SExt;
+
+  auto *ExtInst = CastInst::Create(ExtType, Cond, NewType);
+  ExtInst->insertBefore(SI);
+  SI->setCondition(ExtInst);
+  for (SwitchInst::CaseIt Case : SI->cases()) {
+    APInt NarrowConst = Case.getCaseValue()->getValue();
+    APInt WideConst = (ExtType == Instruction::ZExt) ?
+                      NarrowConst.zext(RegWidth) : NarrowConst.sext(RegWidth);
+    Case.setValue(ConstantInt::get(Context, WideConst));
+  }
+
+  return true;
+}
+
 namespace {
 /// \brief Helper class to promote a scalar operation to a vector one.
 /// This class is used to move downward extractelement transition.
 /// E.g.,
 /// a = vector_op <2 x i32>
 /// b = extractelement <2 x i32> a, i32 0
 /// c = scalar_op b
 /// store c
if (CallInst *CI = dyn_cast<CallInst>(I))
     return optimizeCallInst(CI, ModifiedDT);
 
   if (SelectInst *SI = dyn_cast<SelectInst>(I))
     return optimizeSelectInst(SI);
 
   if (ShuffleVectorInst *SVI = dyn_cast<ShuffleVectorInst>(I))
     return optimizeShuffleVectorInst(SVI);
 
+  if (auto *Switch = dyn_cast<SwitchInst>(I))
+    return optimizeSwitchInst(Switch);
+
   if (isa<ExtractElementInst>(I))
     return optimizeExtractElementInst(I);
 
   return false;
 }
 
 // In this pass we look for GEP and cast instructions that are used
 // across basic blocks and rewrite them to improve basic-block-at-a-time

llvm/trunk/test/Transforms/CodeGenPrepare/widen_switch.ll

+;; PowerPC is arbitralily chosen as a 32/64-bit RISC representative to show the transform in all tests.
+;; x86 is chosen to show that the transform may differ when 8-bit and 16-bit registers are available.
+
+; RUN: opt < %s -codegenprepare -S -mtriple=powerpc64-unknown-unknown | FileCheck %s --check-prefix=PPC --check-prefix=ALL
+; RUN: opt < %s -codegenprepare -S -mtriple=x86_64-unknown-unknown    | FileCheck %s --check-prefix=X86 --check-prefix=ALL
+
+; PPC widens to 32-bit; no change for x86 because 16-bit registers are part of the architecture.
+
+define i32 @widen_switch_i16(i32 %a)  {
+entry:
+  %trunc = trunc i32 %a to i16
+  switch i16 %trunc, label %sw.default [
+    i16 1, label %sw.bb0
+    i16 -1, label %sw.bb1
+  ]
+
+sw.bb0:
+  br label %return
+
+sw.bb1:
+  br label %return
+
+sw.default:
+  br label %return
+
+return:
+  %retval = phi i32 [ -1, %sw.default ], [ 0, %sw.bb0 ], [ 1, %sw.bb1 ]
+  ret i32 %retval
+
+; ALL-LABEL: @widen_switch_i16(
+; PPC:       %0 = zext i16 %trunc to i32
+; PPC-NEXT:  switch i32 %0, label %sw.default [
+; PPC-NEXT:    i32 1, label %return
+; PPC-NEXT:    i32 65535, label %sw.bb1
+; 
+; X86:       %trunc = trunc i32 %a to i16
+; X86-NEXT:  switch i16 %trunc, label %sw.default [
+; X86-NEXT:    i16 1, label %return
+; X86-NEXT:    i16 -1, label %sw.bb1
+}
+
+; Both architectures widen to 32-bit from a smaller, non-native type.
+
+define i32 @widen_switch_i17(i32 %a)  {
+entry:
+  %trunc = trunc i32 %a to i17
+  switch i17 %trunc, label %sw.default [
+    i17 10, label %sw.bb0
+    i17 -1, label %sw.bb1
+  ]
+
+sw.bb0:
+  br label %return
+
+sw.bb1:
+  br label %return
+
+sw.default:
+  br label %return
+
+return:
+  %retval = phi i32 [ -1, %sw.default ], [ 0, %sw.bb0 ], [ 1, %sw.bb1 ]
+  ret i32 %retval
+
+; ALL-LABEL: @widen_switch_i17(
+; ALL:       %0 = zext i17 %trunc to i32
+; ALL-NEXT:  switch i32 %0, label %sw.default [
+; ALL-NEXT:    i32 10, label %return
+; ALL-NEXT:    i32 131071, label %sw.bb1
+}
+
+; If the switch condition is a sign-extended function argument, then the
+; condition and cases should be sign-extended rather than zero-extended
+; because the sign-extension can be optimized away.
+
+define i32 @widen_switch_i16_sext(i2 signext %a)  {
+entry:
+  switch i2 %a, label %sw.default [
+    i2 1, label %sw.bb0
+    i2 -1, label %sw.bb1
+  ]
+
+sw.bb0:
+  br label %return
+
+sw.bb1:
+  br label %return
+
+sw.default:
+  br label %return
+
+return:
+  %retval = phi i32 [ -1, %sw.default ], [ 0, %sw.bb0 ], [ 1, %sw.bb1 ]
+  ret i32 %retval
+
+; ALL-LABEL: @widen_switch_i16_sext(
+; PPC:       %0 = sext i2 %a to i32
+; PPC-NEXT:  switch i32 %0, label %sw.default [
+; PPC-NEXT:    i32 1, label %return
+; PPC-NEXT:    i32 -1, label %sw.bb1
+; 
+; X86:       %0 = sext i2 %a to i8
+; X86-NEXT:  switch i8 %0, label %sw.default [
+; X86-NEXT:    i8 1, label %return
+; X86-NEXT:    i8 -1, label %sw.bb1
+}
+