[RISCV] Eliminate unnecessary masking of promoted shift amounts

[asb]

Summary

SelectionDAGBuilder::visitShift will always zero-extend a shift amount when it is promoted to the ShiftAmountTy. This results in zero-extension (masking) which is unnecessary for RISC-V as the shift operations only read the lower 5 or 6 bits (RV32 or RV64).

I initially proposed adding a getExtendForShiftAmount hook so the shift amount can be any-extended (D52975). @efriedma explained this was unsafe, so I have instead eliminate the unnecessary and operations at instruction selection time in a manner similar to X86InstrCompiler.td.

[efriedma]

LGTM with a couple more tests with an explicit AND (to show the pattern triggers for "and i32 %shamt, 31", but not "and i32 %shamt, 15").

I think you also need patterns for sraw etc., but I guess you'll do that separately?

[asb]

Yes, I'll add the appropraite SLLW/SRAW/SRLW patterns in the RV64I codegen patch. Given this improvement isn't specific to RV64I, I wrote this patch to apply against current HEAD.

[asb]

Updating patch with additional test. Thanks for the review @efriedma.

lib/Target/RISCV/RISCVInstrInfo.td

 // A parameterized register class alternative to i32imm/i64imm from Target.td.
 def ixlenimm : Operand<XLenVT> {
   let ParserMatchClass = ImmXLenAsmOperand<"">;
 }
 
 // Standalone (codegen-only) immleaf patterns.
 def simm32     : ImmLeaf<XLenVT, [{return isInt<32>(Imm);}]>;
 def simm32hi20 : ImmLeaf<XLenVT, [{return isShiftedInt<20, 12>(Imm);}]>;
+// A mask value that won't affect significant shift bits.
+def immshiftxlen : ImmLeaf<XLenVT, [{
+  if (Subtarget->is64Bit())
+    return countTrailingOnes<uint64_t>(Imm) >= 6;
+  return countTrailingOnes<uint64_t>(Imm) >= 5;
+}]>;
 
 // Addressing modes.
 // Necessary because a frameindex can't be matched directly in a pattern.
 def AddrFI : ComplexPattern<iPTR, 1, "SelectAddrFI", [frameindex], []>;
 
 // Extract least significant 12 bits from an immediate value and sign extend
 // them.
 def LO12Sext : SDNodeXForm<imm, [{

 def : PatGprSimm12<add, ADDI>;
 def : PatGprGpr<sub, SUB>;
 def : PatGprGpr<or, OR>;
 def : PatGprSimm12<or, ORI>;
 def : PatGprGpr<and, AND>;
 def : PatGprSimm12<and, ANDI>;
 def : PatGprGpr<xor, XOR>;
 def : PatGprSimm12<xor, XORI>;
-def : PatGprGpr<shl, SLL>;
 def : PatGprUimmLog2XLen<shl, SLLI>;
-def : PatGprGpr<srl, SRL>;
 def : PatGprUimmLog2XLen<srl, SRLI>;
-def : PatGprGpr<sra, SRA>;
 def : PatGprUimmLog2XLen<sra, SRAI>;
 
+// Match both a plain shift and one where the shift amount is masked (this is
+// typically introduced when the legalizer promotes the shift amount and
+// zero-extends it). For RISC-V, the mask is unnecessary as shifts in the base
+// ISA only read the least significant 5 bits (RV32I) or 6 bits (RV64I).
+multiclass VarShiftXLenPat<PatFrag ShiftOp, RVInst Inst> {
+  def : Pat<(ShiftOp GPR:$rs1, GPR:$rs2), (Inst GPR:$rs1, GPR:$rs2)>;
+  def : Pat<(ShiftOp GPR:$rs1, (and GPR:$rs2, immshiftxlen)),
+            (Inst GPR:$rs1, GPR:$rs2)>;
+}
+
+defm : VarShiftXLenPat<shl, SLL>;
+defm : VarShiftXLenPat<srl, SRL>;
+defm : VarShiftXLenPat<sra, SRA>;
+
 /// FrameIndex calculations
 
 def : Pat<(add (i32 AddrFI:$Rs), simm12:$imm12),
           (ADDI (i32 AddrFI:$Rs), simm12:$imm12)>;
 def : Pat<(IsOrAdd (i32 AddrFI:$Rs), simm12:$imm12),
           (ADDI (i32 AddrFI:$Rs), simm12:$imm12)>;
 
 /// Setcc

test/CodeGen/RISCV/alu16.ll

 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
 ; RUN: llc -mtriple=riscv32 -verify-machineinstrs < %s \
 ; RUN:   | FileCheck %s -check-prefix=RV32I
 
 ; These tests are identical to those in alu32.ll but operate on i16. They check
 ; that legalisation of these non-native types doesn't introduce unnecessary
 ; inefficiencies.
 
-; TODO: it's unnecessary to mask (zero-extend) the shift amount.
-
 define i16 @addi(i16 %a) nounwind {
 ; RV32I-LABEL: addi:
 ; RV32I:       # %bb.0:
 ; RV32I-NEXT:    addi a0, a0, 1
 ; RV32I-NEXT:    ret
   %1 = add i16 %a, 1
   ret i16 %1
 }

 ; 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

 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
 ; RUN: llc -mtriple=riscv32 -verify-machineinstrs < %s \
 ; RUN:   | FileCheck %s -check-prefix=RV32I
 
 ; These tests are identical to those in alu32.ll but operate on i8. They check
 ; that legalisation of these non-native types doesn't introduce unnecessary
 ; inefficiencies.
 
-; TODO: it's unnecessary to mask (zero-extend) the shift amount.
-
 define i8 @addi(i8 %a) nounwind {
 ; RV32I-LABEL: addi:
 ; RV32I:       # %bb.0:
 ; RV32I-NEXT:    addi a0, a0, 1
 ; RV32I-NEXT:    ret
   %1 = add i8 %a, 1
   ret i8 %1
 }

 ; 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
 }