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

[Simon Pilgrim via Phabricator via llvm-commits]

RKSimon created this revision.
RKSimon added reviewers: dmgreen, efriedma.
Herald added subscribers: hiraditya, kristof.beyls.
Herald added a project: All.
RKSimon requested review of this revision.
Herald added a project: LLVM.

This came about while I was investigating the lost bit extractions on D125836 <https://reviews.llvm.org/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.


Repository:
  rG LLVM Github Monorepo

https://reviews.llvm.org/D127603

Files:
  llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp


Index: llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp
===================================================================
--- llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp
+++ llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp
@@ -1947,12 +1947,12 @@
   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
@@ -1964,36 +1964,42 @@
   // UBFM Value, ShiftImm, BitWide + SrlImm -1
   //
 
-  if (N->getOpcode() != ISD::SRL)
+  uint64_t SrlImm = 0;
+  if (N->getOpcode() != ISD::SRL || !isIntImmediate(N->getOperand(1), SrlImm))
     return false;
 
-  uint64_t AndMask = 0;
-  if (!isOpcWithIntImmediate(N->getOperand(0).getNode(), ISD::AND, AndMask))
+  // 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;
 
-  Opd0 = N->getOperand(0).getOperand(0);
-
-  uint64_t SrlImm = 0;
-  if (!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)) {
+  // 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) &&
@@ -2008,7 +2014,7 @@
          "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.
@@ -2130,7 +2136,8 @@
                                       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);


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D127603.436230.patch
Type: text/x-patch
Size: 4083 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20220612/84bd10e0/attachment.bin>