[CGP][RISCV] Teach CodeGenPrepare::optimizeSwitchInst to honor isSExtCheaperThanZExt.

[craig.topper]

Summary

This optimization pre-promotes the input and constants for a
switch instruction to a legal type so that all the generated compares
share the same extend. Since RISCV prefers sext for i32 to i64
extends, we should honor that to use sext.w instead of a pair
of shifts.

[frasercrmck]

Seems reasonable to me but I'll let others take a look

[jrtc27]

llvm/lib/CodeGen/CodeGenPrepare.cpp

7018–7020

Do we not want to avoid this in the case this is a function argument that is already being zero-extended? i.e. something like:

Instruction::CastOps ExtType = Instruction::CastOpsEnd;
if (auto *Arg = dyn_cast<Argument>(Cond)) {
  if (Arg->hasSExtAttr())
    ExtType = Instruction::SExt;
  else if (Arg->hasZExtAttr())
    ExtType = Instruction::ZExt;
}
if (ExtType == Instruction::CastOpsEnd) {
  if (TLI->isSExtCheaperThanZExt(OldVT, RegType))
    ExtType = Instruction::SExt;
  else
    ExtType = Instruction::ZExt;
}

(slightly abusing CastOpsEnd as being one past the last valid op, and thus a sentinel, though you could write it in other ways too if you wanted)

[jrtc27]

llvm/lib/CodeGen/CodeGenPrepare.cpp

7018–7020

Or maybe the fact that sext is cheaper means the saving for each case's extending to match outweighs the cost of having to also extend the argument?

[craig.topper]

llvm/lib/CodeGen/CodeGenPrepare.cpp

7018–7020

It's a non-issue for RISCV since I don't think we ever create i32 zext arguments, but it would be more consistent. Could probably just do the isSExtCheaperThanZExt first and then check if it is an Argument after.

[jrtc27]

llvm/lib/CodeGen/CodeGenPrepare.cpp

7018–7020

True, only i16 and i8, and those aren't so cheap to sext. Your suggestion is indeed the cleaner way to express the same thing, not sure how I missed that...

[craig.topper]

Honor zext function attributes over a preference for sext.

llvm/lib/CodeGen/CodeGenPrepare.cpp

bool CodeGenPrepare::tryToSinkFreeOperands(Instruction *I) {
 
   return Changed;
 }
 
 bool CodeGenPrepare::optimizeSwitchInst(SwitchInst *SI) {
   Value *Cond = SI->getCondition();
   Type *OldType = Cond->getType();
   LLVMContext &Context = Cond->getContext();
-  MVT RegType = TLI->getRegisterType(Context, TLI->getValueType(*DL, OldType));
+  EVT OldVT = TLI->getValueType(*DL, OldType);
+  MVT RegType = TLI->getRegisterType(Context, OldVT);
   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.
+  // Extend the switch condition and case constants using the target preferred
+  // extend unless the switch condition is a function argument with an extend
+  // attribute. In that case, we can avoid an unnecessary mask/extension by
+  // matching the argument extension instead.
   Instruction::CastOps ExtType = Instruction::ZExt;
-  if (auto *Arg = dyn_cast<Argument>(Cond))
+  // Some targets prefer SExt over ZExt.
+  if (TLI->isSExtCheaperThanZExt(OldVT, RegType))
+    ExtType = Instruction::SExt;
+
+  if (auto *Arg = dyn_cast<Argument>(Cond)) {
     if (Arg->hasSExtAttr())
       ExtType = Instruction::SExt;
+    if (Arg->hasZExtAttr())
+      ExtType = Instruction::ZExt;
+  }
 
   auto *ExtInst = CastInst::Create(ExtType, Cond, NewType);
   ExtInst->insertBefore(SI);
   ExtInst->setDebugLoc(SI->getDebugLoc());
   SI->setCondition(ExtInst);
   for (auto Case : SI->cases()) {
     APInt NarrowConst = Case.getCaseValue()->getValue();
     APInt WideConst = (ExtType == Instruction::ZExt) ?

llvm/test/CodeGen/RISCV/jumptable.ll

 ; RV32I-MEDIUM-NEXT:    addi a0, zero, 2
 ; RV32I-MEDIUM-NEXT:  .LBB0_9: # %exit
 ; RV32I-MEDIUM-NEXT:    sw a0, 0(a1)
 ; RV32I-MEDIUM-NEXT:  .LBB0_10: # %exit
 ; RV32I-MEDIUM-NEXT:    ret
 ;
 ; RV64I-SMALL-LABEL: below_threshold:
 ; RV64I-SMALL:       # %bb.0: # %entry
-; RV64I-SMALL-NEXT:    slli a0, a0, 32
-; RV64I-SMALL-NEXT:    srli a0, a0, 32
+; RV64I-SMALL-NEXT:    sext.w a0, a0
 ; RV64I-SMALL-NEXT:    addi a2, zero, 2
 ; RV64I-SMALL-NEXT:    blt a2, a0, .LBB0_4
 ; RV64I-SMALL-NEXT:  # %bb.1: # %entry
 ; RV64I-SMALL-NEXT:    addi a2, zero, 1
 ; RV64I-SMALL-NEXT:    beq a0, a2, .LBB0_7
 ; RV64I-SMALL-NEXT:  # %bb.2: # %entry
 ; RV64I-SMALL-NEXT:    addi a2, zero, 2
 ; RV64I-SMALL-NEXT:    bne a0, a2, .LBB0_10

 ; RV64I-SMALL-NEXT:    addi a0, zero, 2
 ; RV64I-SMALL-NEXT:  .LBB0_9: # %exit
 ; RV64I-SMALL-NEXT:    sw a0, 0(a1)
 ; RV64I-SMALL-NEXT:  .LBB0_10: # %exit
 ; RV64I-SMALL-NEXT:    ret
 ;
 ; RV64I-MEDIUM-LABEL: below_threshold:
 ; RV64I-MEDIUM:       # %bb.0: # %entry
-; RV64I-MEDIUM-NEXT:    slli a0, a0, 32
-; RV64I-MEDIUM-NEXT:    srli a0, a0, 32
+; RV64I-MEDIUM-NEXT:    sext.w a0, a0
 ; RV64I-MEDIUM-NEXT:    addi a2, zero, 2
 ; RV64I-MEDIUM-NEXT:    blt a2, a0, .LBB0_4
 ; RV64I-MEDIUM-NEXT:  # %bb.1: # %entry
 ; RV64I-MEDIUM-NEXT:    addi a2, zero, 1
 ; RV64I-MEDIUM-NEXT:    beq a0, a2, .LBB0_7
 ; RV64I-MEDIUM-NEXT:  # %bb.2: # %entry
 ; RV64I-MEDIUM-NEXT:    addi a2, zero, 2
 ; RV64I-MEDIUM-NEXT:    bne a0, a2, .LBB0_10

 ; RV32I-MEDIUM-NEXT:    addi a0, zero, 200
 ; RV32I-MEDIUM-NEXT:  .LBB1_8: # %exit
 ; RV32I-MEDIUM-NEXT:    sw a0, 0(a1)
 ; RV32I-MEDIUM-NEXT:  .LBB1_9: # %exit
 ; RV32I-MEDIUM-NEXT:    ret
 ;
 ; RV64I-SMALL-LABEL: above_threshold:
 ; RV64I-SMALL:       # %bb.0: # %entry
-; RV64I-SMALL-NEXT:    slli a0, a0, 32
-; RV64I-SMALL-NEXT:    srli a0, a0, 32
+; RV64I-SMALL-NEXT:    sext.w a0, a0
 ; RV64I-SMALL-NEXT:    addi a0, a0, -1
 ; RV64I-SMALL-NEXT:    addi a2, zero, 5
 ; RV64I-SMALL-NEXT:    bltu a2, a0, .LBB1_9
 ; RV64I-SMALL-NEXT:  # %bb.1: # %entry
 ; RV64I-SMALL-NEXT:    slli a0, a0, 3
 ; RV64I-SMALL-NEXT:    lui a2, %hi(.LJTI1_0)
 ; RV64I-SMALL-NEXT:    addi a2, a2, %lo(.LJTI1_0)
 ; RV64I-SMALL-NEXT:    add a0, a0, a2

 ; RV64I-SMALL-NEXT:    addi a0, zero, 200
 ; RV64I-SMALL-NEXT:  .LBB1_8: # %exit
 ; RV64I-SMALL-NEXT:    sw a0, 0(a1)
 ; RV64I-SMALL-NEXT:  .LBB1_9: # %exit
 ; RV64I-SMALL-NEXT:    ret
 ;
 ; RV64I-MEDIUM-LABEL: above_threshold:
 ; RV64I-MEDIUM:       # %bb.0: # %entry
-; RV64I-MEDIUM-NEXT:    slli a0, a0, 32
-; RV64I-MEDIUM-NEXT:    srli a0, a0, 32
+; RV64I-MEDIUM-NEXT:    sext.w a0, a0
 ; RV64I-MEDIUM-NEXT:    addi a0, a0, -1
 ; RV64I-MEDIUM-NEXT:    addi a2, zero, 5
 ; RV64I-MEDIUM-NEXT:    bltu a2, a0, .LBB1_9
 ; RV64I-MEDIUM-NEXT:  # %bb.1: # %entry
 ; RV64I-MEDIUM-NEXT:    slli a0, a0, 3
 ; RV64I-MEDIUM-NEXT:  .LBB1_10: # %entry
 ; RV64I-MEDIUM-NEXT:    # Label of block must be emitted
 ; RV64I-MEDIUM-NEXT:    auipc a2, %pcrel_hi(.LJTI1_0)