[llvm] [X86] EltsFromConsecutiveLoads - handle trunc(wideload()) patterns (PR #199371)

[Simon Pilgrim via llvm-commits]

================
@@ -8055,6 +8055,104 @@ static SDValue EltsFromConsecutiveLoads(EVT VT, ArrayRef<SDValue> Elts,
     }
   }
 
+  // STRIDED - element loads at a uniform byte stride larger than the element
+  // size are folded into wide load(s) + vector truncation.
+  // Depth 1 lets the REVERSE block below recurse into us to catch reverse
+  // strides.
+  if (Depth <= 1 && Subtarget.hasAVX2() && !IsConsecutiveLoad &&
+      LoadMask.isAllOnes() && isPowerOf2_32(NumElems) && BaseSizeInBits <= 32 &&
+      VT.getSizeInBits() >= 128) {
+    unsigned WideEltBits = 0;
+    for (unsigned Trial : {16u, 32u, 64u}) {
+      if (Trial <= BaseSizeInBits || Trial % BaseSizeInBits != 0)
+        continue;
+      unsigned LaneStride = Trial / BaseSizeInBits;
+      bool AllMatch = true;
+      for (unsigned k = 1; k < NumElems && AllMatch; ++k) {
+        AllMatch = ByteOffsets[k] == 0 &&
+                   DAG.areNonVolatileConsecutiveLoads(
+                       Loads[k], LDBase, BaseSizeInBytes, k * LaneStride);
+      }
+      if (AllMatch) {
+        WideEltBits = Trial;
+        break;
+      }
+    }
+    if (WideEltBits != 0) {
+      MVT WideEltVT = MVT::getIntegerVT(WideEltBits);
+      MVT SrcEltVT = MVT::getIntegerVT(BaseSizeInBits);
+      // VTRUNC writes the truncated values to the low lanes of an xmm with
+      // zero padding above.
+      MVT TruncDstVT = MVT::getVectorVT(SrcEltVT, 128 / BaseSizeInBits);
+      unsigned TruncDstLanes = TruncDstVT.getVectorNumElements();
+      // Try wider register sizes first.
+      for (unsigned WideRegBits : {512u, 256u, 128u}) {
+        unsigned LanesPerWideLoad = WideRegBits / WideEltBits;
+        if (LanesPerWideLoad < 2 || NumElems % LanesPerWideLoad != 0)
+          continue;
+        if (LanesPerWideLoad > TruncDstLanes)
+          continue; // VTRUNC dest must hold all good lanes of one piece.
+        MVT WideVT = MVT::getVectorVT(WideEltVT, LanesPerWideLoad);
+        if (!TLI.isTypeLegal(WideVT) || !TLI.isTypeLegal(TruncDstVT))
+          continue;
+        unsigned NumWideLoads = NumElems / LanesPerWideLoad;
+        // VTRUNC has no non-AVX-512 lowering so the i16 to i8 form needs BWI.
+        bool Partial = LanesPerWideLoad != TruncDstLanes;
+        if (Partial && (!Subtarget.hasAVX512() ||
+                        (WideEltBits == 16 && !Subtarget.hasBWI())))
+          continue;
+        unsigned BytesPerWideLoad = WideRegBits / 8;
+        auto MMOFlags = LDBase->getMemOperand()->getFlags();
+        SDValue BasePtr = LDBase->getBasePtr();
+        SmallVector<SDValue, 8> Pieces;
+        Pieces.reserve(NumWideLoads);
----------------
RKSimon wrote:

don't bother with reserve - relying on the SmallVector<8> should be fine in most casess

https://github.com/llvm/llvm-project/pull/199371