[llvm] [SelectionDAG] Merge consecutive loads feeding as shuffle operands (PR #207303)

[Ricardo Jesus via llvm-commits]

================
@@ -27007,6 +27007,91 @@ static SDValue combineConcatVectorOfShuffleAndItsOperands(
   return DAG.getVectorShuffle(VT, dl, ShufOps[0], ShufOps[1], Mask);
 }
 
+// concat(shuffle(loadA, loadB, mask0), shuffle(loadA, loadB, mask1))
+// -> shuffle(loadAB, poison, concat(mask0, mask1))
+// only if loadA and loadB can be proven consecutive.
+static SDValue combineConcatVectorOfShuffles(SDNode *N, SelectionDAG &DAG,
+                                             const TargetLowering &TLI) {
+  SDValue A, B;
+  ArrayRef<int> M0, M1;
+  if (!sd_match(N,
+                m_Node(ISD::CONCAT_VECTORS,
+                       m_OneUse(m_Shuffle(m_NUses<2>(m_Value(A)),
+                                          m_NUses<2>(m_Value(B)), m_Mask(M0))),
+                       m_OneUse(m_Shuffle(m_Deferred(A), m_Deferred(B),
+                                          m_Mask(M1))))))
+    return SDValue();
+  auto *L00 = dyn_cast<LoadSDNode>(A.getNode());
+  auto *L01 = dyn_cast<LoadSDNode>(B.getNode());
+  if (!L00 || !L01 || !ISD::isNON_EXTLoad(L00) || !ISD::isNON_EXTLoad(L01))
+    return SDValue();
+
+  // Check if the address spaces of both loads are the same.
+  if (L00->getAddressSpace() != L01->getAddressSpace())
+    return SDValue();
+
+  // Check if the wide load would be faster than the two separate loads.
+  EVT WideVT =
+      L00->getMemoryVT().getDoubleNumVectorElementsVT(*DAG.getContext());
+
+  // Check if the loads are consecutive.
+  LoadSDNode *Base = nullptr;
+  if (DAG.areNonVolatileConsecutiveLoads(
+          L01, L00, L01->getMemoryVT().getStoreSize(), /*Dist=*/1)) {
+    Base = L00;
+  } else if (DAG.areNonVolatileConsecutiveLoads(
+                 L00, L01, L00->getMemoryVT().getStoreSize(), /*Dist=*/1)) {
+    Base = L01;
+  } else {
+    return SDValue(); // not adjacent
+  }
+
+  unsigned Fast = 0;
+  Align NewAlign = Base->getAlign();
+  if (!TLI.allowsMemoryAccess(*DAG.getContext(), DAG.getDataLayout(), WideVT,
+                              Base->getAddressSpace(), NewAlign,
+                              Base->getMemOperand()->getFlags(), &Fast) ||
+      !Fast)
+    return SDValue();
+
+  // Check if this is big endian target. If yes, we need to reverse the wide
+  // load order using bswap, which requires a scalar size that is a multiple
+  // of 16 bits.
+  bool NeedBSwap = DAG.getDataLayout().isBigEndian();
+  if (NeedBSwap && L00->getMemoryVT().getScalarSizeInBits() % 16 != 0)
+    return SDValue();
+
+  // Create a wide load of twice the size of the original load.
+  MachineFunction &MF = DAG.getMachineFunction();
+  MachineMemOperand *WideMMO = MF.getMachineMemOperand(
+      Base->getMemOperand(), /*Offset=*/0, WideVT.getStoreSize());
+  SDValue WideLoad = DAG.getLoad(WideVT, SDLoc(N), Base->getChain(),
+                                 Base->getBasePtr(), WideMMO);
+  SDValue WideMemOp = WideLoad;
+  // Redirect old chain users to the new chain.
+  DAG.makeEquivalentMemoryOrdering(L00, WideMemOp);
+  DAG.makeEquivalentMemoryOrdering(L01, WideMemOp);
+  if (NeedBSwap)
+    WideLoad = DAG.getNode(ISD::BSWAP, SDLoc(N), WideVT, WideLoad);
+
+  // Create a shuffle of the wide load.
+  SmallVector<int, 32> Mask;
+  if (Base == L00) {
+    llvm::append_range(Mask, M0);
+    llvm::append_range(Mask, M1);
+  } else {
+    int Sz = L00->getMemoryVT().getVectorNumElements();
+    for (int I = 0; I < Sz; ++I)
+      Mask.push_back(M0[I] < Sz ? M0[I] + Sz : M0[I] - Sz);
+    for (int I = 0; I < Sz; ++I)
+      Mask.push_back(M1[I] < Sz ? M1[I] + Sz : M1[I] - Sz);
+  }
+  // Create a new shuffle with the new mask.
+  SDValue NewShuffle = DAG.getVectorShuffle(WideVT, SDLoc(N), WideLoad,
+                                            DAG.getPOISON(WideVT), Mask);
+  return NewShuffle;
----------------
rj-jesus wrote:

Can return the new shuffle directly

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