[TargetLowering][RISCV] Introduce getExtendForShiftAmount and implement for RISC-V

[asb]

Summary

As of rL125457, the shift amount operand is always zero-extended when being promoted. This is conservatively safe, but on some targets (e.g. RISC-V) this leads to unnecessary instructions as the native shift operations ignore all but the lower 5/6 bits.

This patch introduces the getExtendForShiftAmount hook which is called from SelectionDAGBuilder::visitShift and implements that hook for RISC-V. The benefit can be seen in the promoted 8 and 16-bit shifts in the test/CodeGen/RISCV/alu{8,16}.ll test cases.

[asb]

Update diff to remove the TODO in alu8.ll and alu16.ll, as this patch addresses that problem.

[efriedma]

You can't ANY_EXTEND here; that will confuse DAGCombine, which will, for example, fold shifts where the shift amount is too large to undef. Other targets normally deal with this at isel time; see, for example, MaskedShiftAmountPats in lib/Target/X86/X86InstrCompiler.td .

[asb]

Thanks for the review, I was hoping there was an opportunity here to avoid the need for backend-specific peepholes. I'll handle in the backend instead. For my own edification, what do you mean by the shift amount being too large to undef?

[efriedma]

I guess that sentence was weird. I meant, if the shift amount is too large, DAGCombine will fold the shift to undef.

[efriedma]

I was hoping there was an opportunity here to avoid the need for backend-specific peepholes

It's possible there's some room for improvement here, to add some target-independent infrastructure for shift instructions that mask the shift amount. But nothing exists at the moment, and I'm not sure what the best approach would be.

[asb]

Thanks for the feedback @efriedma, I've posted a patch that tackles this at instruction selection time: D53224.

include/llvm/CodeGen/TargetLowering.h

public:
 
   /// EVT is not used in-tree, but is used by out-of-tree target.
   /// A documentation for this function would be nice...
   virtual MVT getScalarShiftAmountTy(const DataLayout &, EVT) const;
 
   EVT getShiftAmountTy(EVT LHSTy, const DataLayout &DL,
                        bool LegalTypes = true) const;
 
+  /// Returns how a promoted shift amount value should be extended (ZERO_EXTEND,
+  /// SIGN_EXTEND, or ANY_EXTEND).
+  virtual ISD::NodeType getExtendForShiftAmount(EVT FromShAmtVT,
+                                                EVT ToShAmtVT) const {
+    return ISD::ZERO_EXTEND;
+  }
+
   /// Returns the type to be used for the index operand of:
   /// ISD::INSERT_VECTOR_ELT, ISD::EXTRACT_VECTOR_ELT,
   /// ISD::INSERT_SUBVECTOR, and ISD::EXTRACT_SUBVECTOR
   virtual MVT getVectorIdxTy(const DataLayout &DL) const {
     return getPointerTy(DL);
   }
 
   virtual bool isSelectSupported(SelectSupportKind /*kind*/) const {

lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

void SelectionDAGBuilder::visitBinary(const User &I, unsigned Opcode) {
   SDValue BinNodeValue = DAG.getNode(Opcode, getCurSDLoc(), Op1.getValueType(),
                                      Op1, Op2, Flags);
   setValue(&I, BinNodeValue);
 }
 
 void SelectionDAGBuilder::visitShift(const User &I, unsigned Opcode) {
   SDValue Op1 = getValue(I.getOperand(0));
   SDValue Op2 = getValue(I.getOperand(1));
+  const TargetLowering &TLI = DAG.getTargetLoweringInfo();
 
-  EVT ShiftTy = DAG.getTargetLoweringInfo().getShiftAmountTy(
-      Op1.getValueType(), DAG.getDataLayout());
+  EVT ShiftTy = TLI.getShiftAmountTy(Op1.getValueType(), DAG.getDataLayout());
 
   // Coerce the shift amount to the right type if we can.
   if (!I.getType()->isVectorTy() && Op2.getValueType() != ShiftTy) {
     unsigned ShiftSize = ShiftTy.getSizeInBits();
     unsigned Op2Size = Op2.getValueSizeInBits();
     SDLoc DL = getCurSDLoc();
 
     // If the operand is smaller than the shift count type, promote it.
-    if (ShiftSize > Op2Size)
-      Op2 = DAG.getNode(ISD::ZERO_EXTEND, DL, ShiftTy, Op2);
+    if (ShiftSize > Op2Size) {
+      ISD::NodeType ExtendKind =
+          TLI.getExtendForShiftAmount(Op2.getValueType(), ShiftTy);
+      Op2 = DAG.getNode(ExtendKind, DL, ShiftTy, Op2);
+    }
 
     // If the operand is larger than the shift count type but the shift
     // count type has enough bits to represent any shift value, truncate
     // it now. This is a common case and it exposes the truncate to
     // optimization early.
     else if (ShiftSize >= Log2_32_Ceil(Op2.getValueSizeInBits()))
       Op2 = DAG.getNode(ISD::TRUNCATE, DL, ShiftTy, Op2);
     // Otherwise we'll need to temporarily settle for some other convenient

lib/Target/RISCV/RISCVISelLowering.h

 
 class RISCVTargetLowering : public TargetLowering {
   const RISCVSubtarget &Subtarget;
 
 public:
   explicit RISCVTargetLowering(const TargetMachine &TM,
                                const RISCVSubtarget &STI);
 
+  ISD::NodeType getExtendForShiftAmount(EVT FromShAmtVT,
+                                        EVT ToShAmtVT) const override;
   bool getTgtMemIntrinsic(IntrinsicInfo &Info, const CallInst &I,
                           MachineFunction &MF,
                           unsigned Intrinsic) const override;
   bool isLegalAddressingMode(const DataLayout &DL, const AddrMode &AM, Type *Ty,
                              unsigned AS,
                              Instruction *I = nullptr) const override;
   bool isLegalICmpImmediate(int64_t Imm) const override;
   bool isLegalAddImmediate(int64_t Imm) const override;

lib/Target/RISCV/RISCVISelLowering.cpp

RISCVTargetLowering::RISCVTargetLowering(const TargetMachine &TM,
   unsigned FunctionAlignment = Subtarget.hasStdExtC() ? 1 : 2;
   setMinFunctionAlignment(FunctionAlignment);
   setPrefFunctionAlignment(FunctionAlignment);
 
   // Effectively disable jump table generation.
   setMinimumJumpTableEntries(INT_MAX);
 }
 
+ISD::NodeType
+RISCVTargetLowering::getExtendForShiftAmount(EVT FromShAmtVT,
+                                             EVT ToShAmtVT) const {
+  // RISC-V shift instructions read only the lower 5 or lower 6 bits of the
+  // shift amount. This means ANY_EXTEND is safe as long as the type being
+  // extended is at least that width. Be conservative and assume that there's
+  // still a possibility that the widest supported shift instructions could be
+  // selected.
+  unsigned WidestShAmt = Subtarget.is64Bit() ? 6 : 5;
+  if (FromShAmtVT.getSizeInBits() >= WidestShAmt)
+    return ISD::ANY_EXTEND;
+  return ISD::ZERO_EXTEND;
+}
+
 EVT RISCVTargetLowering::getSetCCResultType(const DataLayout &DL, LLVMContext &,
                                             EVT VT) const {
   if (!VT.isVector())
     return getPointerTy(DL);
   return VT.changeVectorElementTypeToInteger();
 }
 
 bool RISCVTargetLowering::getTgtMemIntrinsic(IntrinsicInfo &Info,

test/CodeGen/RISCV/alu16.ll

 ; RV32I-NEXT:    ret
   %1 = sub i16 %a, %b
   ret i16 %1
 }
 
 define i16 @sll(i16 %a, i16 %b) nounwind {
 ; RV32I-LABEL: sll:
 ; RV32I:       # %bb.0:
-; RV32I-NEXT:    lui a2, 16
-; RV32I-NEXT:    addi a2, a2, -1
-; RV32I-NEXT:    and a1, a1, a2
 ; RV32I-NEXT:    sll a0, a0, a1
 ; RV32I-NEXT:    ret
   %1 = shl i16 %a, %b
   ret i16 %1
 }
 
 define i16 @slt(i16 %a, i16 %b) nounwind {
 ; RV32I-LABEL: slt:
; RV32I-NEXT:    ret
   ret i16 %1
 }
 
 define i16 @srl(i16 %a, i16 %b) nounwind {
 ; RV32I-LABEL: srl:
 ; RV32I:       # %bb.0:
 ; RV32I-NEXT:    lui a2, 16
 ; RV32I-NEXT:    addi a2, a2, -1
-; RV32I-NEXT:    and a1, a1, a2
 ; RV32I-NEXT:    and a0, a0, a2
 ; RV32I-NEXT:    srl a0, a0, a1
 ; RV32I-NEXT:    ret
   %1 = lshr i16 %a, %b
   ret i16 %1
 }
 
 define i16 @sra(i16 %a, i16 %b) nounwind {
 ; RV32I-LABEL: sra:
 ; RV32I:       # %bb.0:
-; RV32I-NEXT:    lui a2, 16
-; RV32I-NEXT:    addi a2, a2, -1
-; RV32I-NEXT:    and a1, a1, a2
 ; RV32I-NEXT:    slli a0, a0, 16
 ; RV32I-NEXT:    srai a0, a0, 16
 ; RV32I-NEXT:    sra a0, a0, a1
 ; RV32I-NEXT:    ret
   %1 = ashr i16 %a, %b
   ret i16 %1
 }
 

test/CodeGen/RISCV/alu8.ll

 ; RV32I-NEXT:    ret
   %1 = sub i8 %a, %b
   ret i8 %1
 }
 
 define i8 @sll(i8 %a, i8 %b) nounwind {
 ; RV32I-LABEL: sll:
 ; RV32I:       # %bb.0:
-; RV32I-NEXT:    andi a1, a1, 255
 ; RV32I-NEXT:    sll a0, a0, a1
 ; RV32I-NEXT:    ret
   %1 = shl i8 %a, %b
   ret i8 %1
 }
 
 define i8 @slt(i8 %a, i8 %b) nounwind {
 ; RV32I-LABEL: slt:

 ; RV32I-NEXT:    ret
   %1 = xor i8 %a, %b
   ret i8 %1
 }
 
 define i8 @srl(i8 %a, i8 %b) nounwind {
 ; RV32I-LABEL: srl:
 ; RV32I:       # %bb.0:
-; RV32I-NEXT:    andi a1, a1, 255
 ; RV32I-NEXT:    andi a0, a0, 255
 ; RV32I-NEXT:    srl a0, a0, a1
 ; RV32I-NEXT:    ret
   %1 = lshr i8 %a, %b
   ret i8 %1
 }
 
 define i8 @sra(i8 %a, i8 %b) nounwind {
 ; RV32I-LABEL: sra:
 ; RV32I:       # %bb.0:
-; RV32I-NEXT:    andi a1, a1, 255
 ; RV32I-NEXT:    slli a0, a0, 24
 ; RV32I-NEXT:    srai a0, a0, 24
 ; RV32I-NEXT:    sra a0, a0, a1
 ; RV32I-NEXT:    ret
   %1 = ashr i8 %a, %b
   ret i8 %1
 }