[RISCV] Select (srl (sext_inreg X, i32), uimm5) to SRAIW if only lower 32 bits are used.

[craig.topper]

Summary

SimplifyDemandedBits can turn srl into sra if the bits being shifted
in aren't demanded. This patch can recover the original sra in some cases.

I've renamed the tablegen class for detecting W users since the "overflowing operator"
term I originally borrowed from Operator.h does not include srl.

[luismarques]

LGTM.

llvm/lib/Target/RISCV/RISCVISelDAGToDAG.cpp

 //===-- RISCVISelDAGToDAG.cpp - A dag to dag inst selector for RISCV ------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //

 // may be able to use a W instruction and CSE with the other instruction if
 // this has happened. We could try to detect that the CSE opportunity exists
 // before doing this, but that would be more complicated.
 // TODO: Does this need to look through AND/OR/XOR to their users to find more
 // opportunities.
 bool RISCVDAGToDAGISel::hasAllNBitUsers(SDNode *Node, unsigned Bits) const {
   assert((Node->getOpcode() == ISD::ADD || Node->getOpcode() == ISD::SUB ||
           Node->getOpcode() == ISD::MUL || Node->getOpcode() == ISD::SHL ||
+          Node->getOpcode() == ISD::SRL ||
           Node->getOpcode() == ISD::SIGN_EXTEND_INREG ||
           isa<ConstantSDNode>(Node)) &&
          "Unexpected opcode");
 
   for (auto UI = Node->use_begin(), UE = Node->use_end(); UI != UE; ++UI) {
     SDNode *User = *UI;
     // Users of this node should have already been instruction selected
     if (!User->isMachineOpcode())

llvm/lib/Target/RISCV/RISCVInstrInfo.td

 // If we're shifting a 32-bit zero extended value left by 0-31 bits, use 2
 // shifts instead of 3. This can occur when unsigned is used to index an array.
 def : Pat<(i64 (shl (and GPR:$rs1, 0xffffffff), uimm5:$shamt)),
           (SRLI (SLLI GPR:$rs1, 32), (ImmSubFrom32 uimm5:$shamt))>;
 }
 
 // PatFrag to allow ADDW/SUBW/MULW/SLLW to be selected from i64 add/sub/mul/shl
 // if only the lower 32 bits of their result is used.
-class overflowingbinopw<SDPatternOperator operator>
+class binop_allwusers<SDPatternOperator operator>
     : PatFrag<(ops node:$lhs, node:$rhs),
               (operator node:$lhs, node:$rhs), [{
   return hasAllWUsers(Node);
 }]>;
 
 def sexti32_allwusers : PatFrag<(ops node:$src),
                                 (sext_inreg node:$src, i32), [{
   return hasAllWUsers(Node);
def : Pat<(sra (sext_inreg GPR:$rs1, i32), uimm5:$shamt),
           (SRAIW GPR:$rs1, uimm5:$shamt)>;
 def : Pat<(i64 (sra (shl GPR:$rs1, (i64 32)), uimm6gt32:$shamt)),
           (SRAIW GPR:$rs1, (ImmSub32 uimm6gt32:$shamt))>;
 
 def : PatGprGpr<shiftopw<riscv_sllw>, SLLW>;
 def : PatGprGpr<shiftopw<riscv_srlw>, SRLW>;
 def : PatGprGpr<shiftopw<riscv_sraw>, SRAW>;
 
-// Select W instructions without sext_inreg if only the lower 32 bits of the
-// result are used.
-def : PatGprGpr<overflowingbinopw<add>, ADDW>;
-def : PatGprSimm12<overflowingbinopw<add>, ADDIW>;
-def : PatGprGpr<overflowingbinopw<sub>, SUBW>;
-def : PatGprImm<overflowingbinopw<shl>, SLLIW, uimm5>;
+// Select W instructions if only the lower 32 bits of the result are used.
+def : PatGprGpr<binop_allwusers<add>, ADDW>;
+def : PatGprSimm12<binop_allwusers<add>, ADDIW>;
+def : PatGprGpr<binop_allwusers<sub>, SUBW>;
+def : PatGprImm<binop_allwusers<shl>, SLLIW, uimm5>;
+
+// If this is a shr of a value sign extended from i32, and all the users only
+// use the lower 32 bits, we can use an sraiw to remove the sext_inreg. This
+// occurs because SimplifyDemandedBits prefers srl over sra.
+def : Pat<(binop_allwusers<srl> (sext_inreg GPR:$rs1, i32), uimm5:$shamt),
+          (SRAIW GPR:$rs1, uimm5:$shamt)>;
 
 /// Loads
 
 defm : LdPat<sextloadi32, LW, i64>;
 defm : LdPat<extloadi32, LW, i64>;
 defm : LdPat<zextloadi32, LWU, i64>;
 defm : LdPat<load, LD, i64>;
 

 
 /// Simple optimization
 def : Pat<(add GPR:$rs1, (AddiPair:$rs2)),
           (ADDI (ADDI GPR:$rs1, (AddiPairImmB AddiPair:$rs2)),
                 (AddiPairImmA GPR:$rs2))>;
 
 let Predicates = [IsRV64] in {
 // Select W instructions if only the lower 32-bits of the result are used.
-def : Pat<(overflowingbinopw<add> GPR:$rs1, (AddiPair:$rs2)),
+def : Pat<(binop_allwusers<add> GPR:$rs1, (AddiPair:$rs2)),
           (ADDIW (ADDIW GPR:$rs1, (AddiPairImmB AddiPair:$rs2)),
                  (AddiPairImmA AddiPair:$rs2))>;
 }
 
 //===----------------------------------------------------------------------===//
 // Standard extensions
 //===----------------------------------------------------------------------===//
 
 include "RISCVInstrInfoM.td"
 include "RISCVInstrInfoA.td"
 include "RISCVInstrInfoF.td"
 include "RISCVInstrInfoD.td"
 include "RISCVInstrInfoC.td"
 include "RISCVInstrInfoB.td"
 include "RISCVInstrInfoV.td"
 include "RISCVInstrInfoZfh.td"

llvm/lib/Target/RISCV/RISCVInstrInfoM.td

 def : PatGprGpr<riscv_mulhsu, MULHSU>;
 def : PatGprGpr<sdiv, DIV>;
 def : PatGprGpr<udiv, DIVU>;
 def : PatGprGpr<srem, REM>;
 def : PatGprGpr<urem, REMU>;
 } // Predicates = [HasStdExtM]
 
 let Predicates = [HasStdExtM, IsRV64] in {
-// Select W instructions without sext_inreg if only the lower 32-bits of the
-// result are used.
-def : PatGprGpr<overflowingbinopw<mul>, MULW>;
+// Select W instructions if only the lower 32-bits of the result are used.
+def : PatGprGpr<binop_allwusers<mul>, MULW>;
 
 def : PatGprGpr<riscv_divw, DIVW>;
 def : PatGprGpr<riscv_divuw, DIVUW>;
 def : PatGprGpr<riscv_remuw, REMUW>;
 
 // Handle the specific cases where using DIVU/REMU would be correct and result
 // in fewer instructions than emitting DIVUW/REMUW then zero-extending the
 // result.

llvm/test/CodeGen/RISCV/rv64i-exhaustive-w-insts.ll

 ; RV64I:       # %bb.0:
 ; RV64I-NEXT:    sext.w a0, a0
 ; RV64I-NEXT:    slli a0, a0, 25
 ; RV64I-NEXT:    srli a0, a0, 32
 ; RV64I-NEXT:    ret
 ;
 ; RV64ZBA-LABEL: zext_sraiw_aext:
 ; RV64ZBA:       # %bb.0:
-; RV64ZBA-NEXT:    sext.w a0, a0
-; RV64ZBA-NEXT:    srli a0, a0, 7
+; RV64ZBA-NEXT:    sraiw a0, a0, 7
 ; RV64ZBA-NEXT:    zext.w a0, a0
 ; RV64ZBA-NEXT:    ret
   %1 = ashr i32 %a, 7
   ret i32 %1
 }
 
 define zeroext i32 @zext_sraiw_sext(i32 signext %a) nounwind {
 ; RV64I-LABEL: zext_sraiw_sext:

 ; RV64I:       # %bb.0:
 ; RV64I-NEXT:    sext.w a0, a0
 ; RV64I-NEXT:    slli a0, a0, 23
 ; RV64I-NEXT:    srli a0, a0, 32
 ; RV64I-NEXT:    ret
 ;
 ; RV64ZBA-LABEL: zext_sraiw_zext:
 ; RV64ZBA:       # %bb.0:
-; RV64ZBA-NEXT:    sext.w a0, a0
-; RV64ZBA-NEXT:    srli a0, a0, 9
+; RV64ZBA-NEXT:    sraiw a0, a0, 9
 ; RV64ZBA-NEXT:    zext.w a0, a0
 ; RV64ZBA-NEXT:    ret
   %1 = ashr i32 %a, 9
   ret i32 %1
 }

llvm/test/CodeGen/RISCV/srem-lkk.ll

 ; RV64IM:       # %bb.0:
 ; RV64IM-NEXT:    sext.w a1, a0
 ; RV64IM-NEXT:    lui a2, 706409
 ; RV64IM-NEXT:    addiw a2, a2, 389
 ; RV64IM-NEXT:    mul a1, a1, a2
 ; RV64IM-NEXT:    srli a1, a1, 32
 ; RV64IM-NEXT:    addw a1, a1, a0
 ; RV64IM-NEXT:    srliw a2, a1, 31
-; RV64IM-NEXT:    srli a1, a1, 6
+; RV64IM-NEXT:    sraiw a1, a1, 6
 ; RV64IM-NEXT:    addw a1, a1, a2
 ; RV64IM-NEXT:    addi a2, zero, 95
 ; RV64IM-NEXT:    mulw a1, a1, a2
 ; RV64IM-NEXT:    subw a0, a0, a1
 ; RV64IM-NEXT:    ret
   %1 = srem i32 %x, 95
   ret i32 %1
 }