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

[Sushant Gokhale via llvm-commits]

================
@@ -27007,6 +27007,88 @@ 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_Shuffle(m_Value(A), m_Value(B), m_Mask(M0)),
+                          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());
+
+  unsigned Fast = 0;
+  Align NewAlign = L00->getAlign();
+  if (!TLI.allowsMemoryAccess(*DAG.getContext(), DAG.getDataLayout(), WideVT,
+                              L00->getAddressSpace(), NewAlign,
+                              L00->getMemOperand()->getFlags(), &Fast) ||
+      !Fast)
+    return SDValue();
+
+  // 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
+  }
+
+  // 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();
----------------
sushgokh wrote:

Just used it a precaution. Can drop it if you suggest.

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