[AArch64] isSeveralBitsExtractOpFromShr - match UBFM patterns with value tracking (RFC)

[RKSimon]

Summary

This came about while I was investigating the lost bit extractions on D125836 - it didn't help but is a more general implementation than just matching AND masks.

With a bit more work I'm wondering how many of the 'isBitfieldExtractOp' patterns could actually be performed with just value tracking and SimplifyMultipleUseDemandedBits.

[efriedma]

llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp

1958

The pattern we're looking for is SRL+AND. The modified version of the function doesn't explicitly check for the "AND" operation at all. I guess you can describe "AND" as "an operation with known zero bits in the result, and can be eliminated by SimplifyMultipleUseDemandedBits", but that seems like an awfully confusing description.

How does this change actually generalize the matching in practice? Are we just looking for different AND masks, or can we actually match operations which aren't AND?

Maybe it would make the result easier to read if we extract the code that checks for the masking operation into a separate "MatchMaskOperation" function?

[RKSimon]

llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp

1958

Yes its most likely only ever going to match AND masks.

This was an experiment as a lot of recent SimplifyDemandedBits work (e.g. D125836) has managed to break the fragile bit extract patterns, and I was curious how general value tracking and demanded bits handling would work - although we're limited on how much we can manipulate the DAG at this late stage. I was curious what others thought, so I put up this RFC.

Maybe the aarch64 backend just needs additional canonicalization for some UBFX/SBFX patterns?

[efriedma]

llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp

1958

DAGCombine likes to manipulate basic arithmetic operations, so it's sort of hard to canonicalize them... but maybe some sort of canonicalization would be helpful. We do already have AArch64TargetLowering::isDesirableToCommuteWithShift.

I haven't really thought about the ideal canonical form.

[RKSimon]

D125836 tweaks AArch64TargetLowering::shouldFoldConstantShiftPairToMask to ensure we don't mess with srl(shl(x,c1),c2) if it matches a UBFX pattern.

Abandoning this for now, but we could revive this in the future if it ever becomes useful (I also noticed that ARM has some similar matching). Bit extraction instructions are common enough that a generic opcode could even be considered....

llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp

static bool isBitfieldExtractOpFromSExtInReg(SDNode *N, unsigned &Opc,
 
   Opc = (VT == MVT::i32) ? AArch64::SBFMWri : AArch64::SBFMXri;
   Opd0 = Op.getOperand(0);
   Immr = ShiftImm;
   Imms = ShiftImm + Width - 1;
   return true;
 }
 
-static bool isSeveralBitsExtractOpFromShr(SDNode *N, unsigned &Opc,
-                                          SDValue &Opd0, unsigned &LSB,
-                                          unsigned &MSB) {
+static bool isSeveralBitsExtractOpFromShr(SelectionDAG *CurDAG, SDNode *N,
+                                          unsigned &Opc, SDValue &Opd0,
+                                          unsigned &LSB, unsigned &MSB) {
   // We are looking for the following pattern which basically extracts several
   // continuous bits from the source value and places it from the LSB of the
-  // destination value, all other bits of the destination value or set to zero:
+  // destination value, all other bits of the destination value are set to zero:
   //
   // Value2 = AND Value, MaskImm
   // SRL Value2, ShiftImm
   //
   // with MaskImm >> ShiftImm to search for the bit width.
   //
   // This gets selected into a single UBFM:
   //
   // UBFM Value, ShiftImm, BitWide + SrlImm -1
   //
 
-  if (N->getOpcode() != ISD::SRL)
-    return false;
-
-  uint64_t AndMask = 0;
-  if (!isOpcWithIntImmediate(N->getOperand(0).getNode(), ISD::AND, AndMask))
-    return false;
-
-  Opd0 = N->getOperand(0).getOperand(0);
-
   uint64_t SrlImm = 0;
-  if (!isIntImmediate(N->getOperand(1), SrlImm))
+  if (N->getOpcode() != ISD::SRL || !isIntImmediate(N->getOperand(1), SrlImm))
     return false;
 
-  // Check whether we really have several bits extract here.
-  unsigned BitWide = 64 - countLeadingOnes(~(AndMask >> SrlImm));
-  if (BitWide && isMask_64(AndMask >> SrlImm)) {
+  // Determine the number of known active bits that are being shifted.
+  SDValue ShiftSrc = N->getOperand(0);
+  KnownBits Known = CurDAG->computeKnownBits(ShiftSrc);
+  unsigned NumBits = Known.getBitWidth();
+  unsigned MinLZ = Known.countMinLeadingZeros(); 
+  unsigned ActiveBits = NumBits - MinLZ;
+  if (ActiveBits <= SrlImm)
+    return false;
+
+  // UBFM will implicitly zero-extend the result, so just based on the active
+  // bits we actually want to extract, see if we can peek through to an inner
+  // source and avoid the mask entirely.
+  const TargetLowering &TLI = CurDAG->getTargetLoweringInfo();
+  APInt DemandedBits = APInt::getBitsSet(NumBits, SrlImm, ActiveBits);
+  if (SDValue Src = TLI.SimplifyMultipleUseDemandedBits(ShiftSrc, DemandedBits,
+                                                        *CurDAG)) {
     if (N->getValueType(0) == MVT::i32)
       Opc = AArch64::UBFMWri;
     else
       Opc = AArch64::UBFMXri;
 
+    Opd0 = Src;
     LSB = SrlImm;
-    MSB = BitWide + SrlImm - 1;
+    MSB = ActiveBits - 1;
     return true;
   }
 
   return false;
 }
 
-static bool isBitfieldExtractOpFromShr(SDNode *N, unsigned &Opc, SDValue &Opd0,
+static bool isBitfieldExtractOpFromShr(SelectionDAG *CurDAG, SDNode *N,
+                                       unsigned &Opc, SDValue &Opd0,
                                        unsigned &Immr, unsigned &Imms,
                                        bool BiggerPattern) {
   assert((N->getOpcode() == ISD::SRA || N->getOpcode() == ISD::SRL) &&
          "N must be a SHR/SRA operation to call this function");
 
   EVT VT = N->getValueType(0);
 
   // Here we can test the type of VT and return false when the type does not
   // match, but since it is done prior to that call in the current context
   // we turned that into an assert to avoid redundant code.
   assert((VT == MVT::i32 || VT == MVT::i64) &&
          "Type checking must have been done before calling this function");
 
   // Check for AND + SRL doing several bits extract.
-  if (isSeveralBitsExtractOpFromShr(N, Opc, Opd0, Immr, Imms))
+  if (isSeveralBitsExtractOpFromShr(CurDAG, N, Opc, Opd0, Immr, Imms))
     return true;
 
   // We're looking for a shift of a shift.
   uint64_t ShlImm = 0;
   uint64_t TruncBits = 0;
   if (isOpcWithIntImmediate(N->getOperand(0).getNode(), ISD::SHL, ShlImm)) {
     Opd0 = N->getOperand(0).getOperand(0);
   } else if (VT == MVT::i32 && N->getOpcode() == ISD::SRL &&
default:
     if (!N->isMachineOpcode())
       return false;
     break;
   case ISD::AND:
     return isBitfieldExtractOpFromAnd(CurDAG, N, Opc, Opd0, Immr, Imms,
                                       NumberOfIgnoredLowBits, BiggerPattern);
   case ISD::SRL:
   case ISD::SRA:
-    return isBitfieldExtractOpFromShr(N, Opc, Opd0, Immr, Imms, BiggerPattern);
+    return isBitfieldExtractOpFromShr(CurDAG, N, Opc, Opd0, Immr, Imms,
+                                      BiggerPattern);
 
   case ISD::SIGN_EXTEND_INREG:
     return isBitfieldExtractOpFromSExtInReg(N, Opc, Opd0, Immr, Imms);
   }
 
   unsigned NOpc = N->getMachineOpcode();
   switch (NOpc) {
   default: