[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);
+        for (unsigned k = 0; k < NumWideLoads; ++k) {
+          unsigned Offset = k * BytesPerWideLoad;
+          SDValue Ptr = k == 0 ? BasePtr
+                               : DAG.getMemBasePlusOffset(
+                                     BasePtr, TypeSize::getFixed(Offset), DL);
+          SDValue Ld =
+              DAG.getLoad(WideVT, DL, LDBase->getChain(), Ptr,
+                          LDBase->getPointerInfo().getWithOffset(Offset),
+                          k == 0 ? LDBase->getBaseAlign() : Align(1), MMOFlags);
+          for (auto *LD : Loads)
+            if (LD)
+              DAG.makeEquivalentMemoryOrdering(LD, Ld);
+          unsigned TruncOp = Partial ? X86ISD::VTRUNC : ISD::TRUNCATE;
+          Pieces.push_back(DAG.getNode(TruncOp, DL, TruncDstVT, Ld));
+        }
+        // Pairwise shuffle the low halves until each piece is full.
+        if (Partial) {
+          unsigned GoodLanes = LanesPerWideLoad;
+          while (Pieces.size() > 1 && GoodLanes < TruncDstLanes) {
+            SmallVector<SDValue, 8> Next;
+            SmallVector<int, 16> Mask(TruncDstLanes, -1);
+            for (unsigned i = 0; i < GoodLanes; ++i) {
+              Mask[i] = i;
+              Mask[GoodLanes + i] = TruncDstLanes + i;
+            }
+            for (unsigned j = 0; j + 1 < Pieces.size(); j += 2)
+              Next.push_back(DAG.getVectorShuffle(TruncDstVT, DL, Pieces[j],
+                                                  Pieces[j + 1], Mask));
+            Pieces = std::move(Next);
+            GoodLanes *= 2;
+          }
+        }
+        SDValue Result;
+        if (Pieces.size() == 1) {
+          Result = Pieces[0];
----------------
RKSimon wrote:

```
if (Pieces.size() == 1)
  return DAG.getBitcast(VT, Pieces[0]);

MVT ConcatVT = MVT::getVectorVT(SrcEltVT, Pieces.size() * TruncDstLanes);
if (!TLI.isTypeLegal(ConcatVT))
  continue;

Result = DAG.getNode(ISD::CONCAT_VECTORS, DL, ConcatVT, Pieces);
return DAG.getBitcast(VT, Result);
```

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