[RISCV] Teach targetShrinkDemandedConstant to preserve (and X, 0xffffffff).

[craig.topper]

Summary

We look for this pattern frequently in isel patterns so its a
good idea to try to preserve it.

This also let's us remove our special isel handling for srliw
and use a direct pattern match of (srl (and X, 0xffffffff), C)
since no bits will be removed from the and mask.

[frasercrmck]

Are there situations in which preserving the and benefits codegen, aside from allowing us to remove the SRLIWPat?

[craig.topper]

In D99042#2641799, @frasercrmck wrote:

Are there situations in which preserving the and benefits codegen, aside from allowing us to remove the SRLIWPat?

https://godbolt.org/z/GnsMvavrv though arguably from that example we should be preserving any trailing 1s mask and teach isel to use SLLI+SRLI for any trailing mask not just 0xffffffff.

[craig.topper]

Here's another https://godbolt.org/z/11s4Mva9v

[frasercrmck]

In D99042#2641992, @craig.topper wrote:

Here's another https://godbolt.org/z/11s4Mva9v

Thanks. Perhaps we can pre-commit those test cases to better show the effect of this optimization?

[craig.topper]

Rebase on top of new tests

[frasercrmck]

LGTM, cheers.

llvm/lib/Target/RISCV/RISCVISelDAGToDAG.h

public:
   }
   bool selectShiftMask32(SDValue N, SDValue &ShAmt) {
     return selectShiftMask(N, 32, ShAmt);
   }
 
   bool selectSExti32(SDValue N, SDValue &Val);
   bool selectZExti32(SDValue N, SDValue &Val);
 
-  bool MatchSRLIW(SDNode *N) const;
   bool MatchSLLIUW(SDNode *N) const;
 
   bool selectVLOp(SDValue N, SDValue &VL);
 
   bool selectVSplat(SDValue N, SDValue &SplatVal);
   bool selectVSplatSimm5(SDValue N, SDValue &SplatVal);
   bool selectVSplatUimm5(SDValue N, SDValue &SplatVal);
 

llvm/lib/Target/RISCV/RISCVISelDAGToDAG.cpp

if (N.getOpcode() == ISD::AssertZext &&
       cast<VTSDNode>(N->getOperand(1))->getVT().bitsLE(MVT::i32)) {
     Val = N;
     return true;
   }
 
   return false;
 }
 
-// Match (srl (and val, mask), imm) where the result would be a
-// zero-extended 32-bit integer. i.e. the mask is 0xffffffff or the result
-// is equivalent to this (SimplifyDemandedBits may have removed lower bits
-// from the mask that aren't necessary due to the right-shifting).
-bool RISCVDAGToDAGISel::MatchSRLIW(SDNode *N) const {
-  assert(N->getOpcode() == ISD::SRL);
-  assert(N->getOperand(0).getOpcode() == ISD::AND);
-  assert(isa<ConstantSDNode>(N->getOperand(1)));
-  assert(isa<ConstantSDNode>(N->getOperand(0).getOperand(1)));
-
-  // The IsRV64 predicate is checked after PatFrag predicates so we can get
-  // here even on RV32.
-  if (!Subtarget->is64Bit())
-    return false;
-
-  SDValue And = N->getOperand(0);
-  uint64_t ShAmt = N->getConstantOperandVal(1);
-  uint64_t Mask = And.getConstantOperandVal(1);
-  return (Mask | maskTrailingOnes<uint64_t>(ShAmt)) == 0xffffffff;
-}
-
 // Check that it is a SLLIUW (Shift Logical Left Immediate Unsigned i32
 // on RV64).
 // SLLIUW is the same as SLLI except for the fact that it clears the bits
 // XLEN-1:32 of the input RS1 before shifting.
 // A PatFrag has already checked that it has the right structure:
 //
 //  (AND (SHL RS1, VC2), VC1)
 //

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

bool RISCVTargetLowering::targetShrinkDemandedConstant(
   if (!C)
     return false;
 
   const APInt &Mask = C->getAPIntValue();
 
   // Clear all non-demanded bits initially.
   APInt ShrunkMask = Mask & DemandedBits;
 
+  // Try to make a smaller immediate by setting undemanded bits.
+
+  APInt ExpandedMask = Mask | ~DemandedBits;
+
+  auto IsLegalMask = [ShrunkMask, ExpandedMask](const APInt &Mask) -> bool {
+    return ShrunkMask.isSubsetOf(Mask) && Mask.isSubsetOf(ExpandedMask);
+  };
+  auto UseMask = [Mask, Op, VT, &TLO](const APInt &NewMask) -> bool {
+    if (NewMask == Mask)
+      return true;
+    SDLoc DL(Op);
+    SDValue NewC = TLO.DAG.getConstant(NewMask, DL, VT);
+    SDValue NewOp = TLO.DAG.getNode(ISD::AND, DL, VT, Op.getOperand(0), NewC);
+    return TLO.CombineTo(Op, NewOp);
+  };
+
   // If the shrunk mask fits in sign extended 12 bits, let the target
   // independent code apply it.
   if (ShrunkMask.isSignedIntN(12))
     return false;
 
-  // Try to make a smaller immediate by setting undemanded bits.
+  // Try to preserve (and X, 0xffffffff), the (zext_inreg X, i32) pattern.
+  if (VT == MVT::i64) {
+    APInt NewMask = APInt(64, 0xffffffff);
+    if (IsLegalMask(NewMask))
+      return UseMask(NewMask);
+  }
 
-  // We need to be able to make a negative number through a combination of mask
-  // and undemanded bits.
-  APInt ExpandedMask = Mask | ~DemandedBits;
+  // For the remaining optimizations, we need to be able to make a negative
+  // number through a combination of mask and undemanded bits.
   if (!ExpandedMask.isNegative())
     return false;
 
   // What is the fewest number of bits we need to represent the negative number.
   unsigned MinSignedBits = ExpandedMask.getMinSignedBits();
 
   // Try to make a 12 bit negative immediate. If that fails try to make a 32
   // bit negative immediate unless the shrunk immediate already fits in 32 bits.
   APInt NewMask = ShrunkMask;
   if (MinSignedBits <= 12)
     NewMask.setBitsFrom(11);
   else if (MinSignedBits <= 32 && !ShrunkMask.isSignedIntN(32))
     NewMask.setBitsFrom(31);
   else
     return false;
 
   // Sanity check that our new mask is a subset of the demanded mask.
-  assert(NewMask.isSubsetOf(ExpandedMask));
-
-  // If we aren't changing the mask, just return true to keep it and prevent
-  // the caller from optimizing.
-  if (NewMask == Mask)
-    return true;
-
-  // Replace the constant with the new mask.
-  SDLoc DL(Op);
-  SDValue NewC = TLO.DAG.getConstant(NewMask, DL, VT);
-  SDValue NewOp = TLO.DAG.getNode(ISD::AND, DL, VT, Op.getOperand(0), NewC);
-  return TLO.CombineTo(Op, NewOp);
+  assert(IsLegalMask(NewMask));
+  return UseMask(NewMask);
 }
 
 void RISCVTargetLowering::computeKnownBitsForTargetNode(const SDValue Op,
                                                         KnownBits &Known,
                                                         const APInt &DemandedElts,
                                                         const SelectionDAG &DAG,
                                                         unsigned Depth) const {
   unsigned BitWidth = Known.getBitWidth();

llvm/lib/Target/RISCV/RISCVInstrInfo.td

def assertsexti32 : PatFrag<(ops node:$src), (assertsext node:$src), [{
   return cast<VTSDNode>(N->getOperand(1))->getVT().bitsLE(MVT::i32);
 }]>;
 def sexti32 : ComplexPattern<i64, 1, "selectSExti32">;
 def assertzexti32 : PatFrag<(ops node:$src), (assertzext node:$src), [{
   return cast<VTSDNode>(N->getOperand(1))->getVT().bitsLE(MVT::i32);
 }]>;
 def zexti32 : ComplexPattern<i64, 1, "selectZExti32">;
 
-def SRLIWPat : PatFrag<(ops node:$A, node:$B),
-                       (srl (and node:$A, imm), node:$B), [{
-  return MatchSRLIW(N);
-}]>;
-
 // Check that it is a SLLIUW (Shift Logical Left Immediate Unsigned i32
 // on RV64). Also used to optimize the same sequence without SLLIUW.
 def SLLIUWPat : PatFrag<(ops node:$A, node:$B),
                         (and (shl node:$A, node:$B), imm), [{
   return MatchSLLIUW(N);
 }]>;
 
 /// Simple arithmetic operations

 def : Pat<(sext_inreg (add GPR:$rs1, GPR:$rs2), i32),
           (ADDW GPR:$rs1, GPR:$rs2)>;
 def : Pat<(sext_inreg (add GPR:$rs1, simm12:$imm12), i32),
           (ADDIW GPR:$rs1, simm12:$imm12)>;
 def : Pat<(sext_inreg (sub GPR:$rs1, GPR:$rs2), i32),
           (SUBW GPR:$rs1, GPR:$rs2)>;
 def : Pat<(sext_inreg (shl GPR:$rs1, uimm5:$shamt), i32),
           (SLLIW GPR:$rs1, uimm5:$shamt)>;
-def : Pat<(i64 (SRLIWPat GPR:$rs1, uimm5:$shamt)),
+def : Pat<(i64 (srl (and GPR:$rs1, 0xffffffff), uimm5:$shamt)),
           (SRLIW GPR:$rs1, uimm5:$shamt)>;
 def : Pat<(i64 (srl (shl GPR:$rs1, (i64 32)), uimm6gt32:$shamt)),
           (SRLIW GPR:$rs1, (ImmSub32 uimm6gt32:$shamt))>;
 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))>;
 

llvm/lib/Target/RISCV/RISCVInstrInfoB.td

 } // Predicates = [HasStdExtZbb, IsRV64]
 
 let Predicates = [HasStdExtZbp, IsRV64] in {
 def : Pat<(i64 (sext_inreg (or (shl GPR:$rs2, (i64 16)),
                                (and GPR:$rs1, 0x000000000000FFFF)),
                            i32)),
           (PACKW GPR:$rs1, GPR:$rs2)>;
 def : Pat<(i64 (or (and (assertsexti32 GPR:$rs2), 0xFFFFFFFFFFFF0000),
-                   (SRLIWPat GPR:$rs1, (i64 16)))),
+                   (srl (and GPR:$rs1, 0xFFFFFFFF), (i64 16)))),
           (PACKUW GPR:$rs1, GPR:$rs2)>;
 } // Predicates = [HasStdExtZbp, IsRV64]

llvm/test/CodeGen/RISCV/alu32.ll

 ; RV64I-LABEL: srli:
 ; RV64I:       # %bb.0:
 ; RV64I-NEXT:    srliw a0, a0, 8
 ; RV64I-NEXT:    ret
   %1 = lshr i32 %a, 8
   ret i32 %1
 }
 
-; FIXME: This should use srliw on RV64, but SimplifyDemandedBits breaks the
-; (and X, 0xffffffff) that type legalization inserts.
+; This makes sure SimplifyDemandedBits doesn't prevent us from matching SRLIW
+; on RV64.
 define i32 @srli_demandedbits(i32 %0) {
 ; RV32I-LABEL: srli_demandedbits:
 ; RV32I:       # %bb.0:
 ; RV32I-NEXT:    srli a0, a0, 3
 ; RV32I-NEXT:    ori a0, a0, 1
 ; RV32I-NEXT:    ret
 ;
 ; RV64I-LABEL: srli_demandedbits:
 ; RV64I:       # %bb.0:
-; RV64I-NEXT:    addi a1, zero, 1
-; RV64I-NEXT:    slli a1, a1, 32
-; RV64I-NEXT:    addi a1, a1, -16
-; RV64I-NEXT:    and a0, a0, a1
-; RV64I-NEXT:    srli a0, a0, 3
+; RV64I-NEXT:    srliw a0, a0, 3
 ; RV64I-NEXT:    ori a0, a0, 1
 ; RV64I-NEXT:    ret
   %2 = lshr i32 %0, 3
   %3 = or i32 %2, 1
   ret i32 %3
 }
 
 define i32 @srai(i32 %a) nounwind {

llvm/test/CodeGen/RISCV/rv64zba.ll

 ; RV64IBA-LABEL: zextw_i64:
 ; RV64IBA:       # %bb.0:
 ; RV64IBA-NEXT:    zext.w a0, a0
 ; RV64IBA-NEXT:    ret
   %and = and i64 %a, 4294967295
   ret i64 %and
 }
 
-; FIXME: This can use zext.w, but we need targetShrinkDemandedConstant to
-; to adjust the immediate.
+; This makes sure targetShrinkDemandedConstant changes the and immmediate to
+; allow zext.w or slli+srli.
 define i64 @zextw_demandedbits_i64(i64 %0) {
 ; RV64I-LABEL: zextw_demandedbits_i64:
 ; RV64I:       # %bb.0:
-; RV64I-NEXT:    addi a1, zero, 1
-; RV64I-NEXT:    slli a1, a1, 32
-; RV64I-NEXT:    addi a1, a1, -2
-; RV64I-NEXT:    and a0, a0, a1
 ; RV64I-NEXT:    ori a0, a0, 1
+; RV64I-NEXT:    slli a0, a0, 32
+; RV64I-NEXT:    srli a0, a0, 32
 ; RV64I-NEXT:    ret
 ;
 ; RV64IB-LABEL: zextw_demandedbits_i64:
 ; RV64IB:       # %bb.0:
-; RV64IB-NEXT:    addi a1, zero, 1
-; RV64IB-NEXT:    slli a1, a1, 32
-; RV64IB-NEXT:    addi a1, a1, -2
-; RV64IB-NEXT:    and a0, a0, a1
 ; RV64IB-NEXT:    ori a0, a0, 1
+; RV64IB-NEXT:    zext.w a0, a0
 ; RV64IB-NEXT:    ret
 ;
 ; RV64IBA-LABEL: zextw_demandedbits_i64:
 ; RV64IBA:       # %bb.0:
-; RV64IBA-NEXT:    addi a1, zero, 1
-; RV64IBA-NEXT:    slli a1, a1, 32
-; RV64IBA-NEXT:    addi a1, a1, -2
-; RV64IBA-NEXT:    and a0, a0, a1
 ; RV64IBA-NEXT:    ori a0, a0, 1
+; RV64IBA-NEXT:    zext.w a0, a0
 ; RV64IBA-NEXT:    ret
   %2 = and i64 %0, 4294967294
   %3 = or i64 %2, 1
   ret i64 %3
 }
 
 define signext i16 @sh1add(i64 %0, i16* %1) {
 ; RV64I-LABEL: sh1add: