[llvm] [VectorCombine] Combine adjacent loads feeding as shuffle operands (PR #213007)

[Sushant Gokhale via llvm-commits]

================
@@ -6800,6 +6804,211 @@ bool VectorCombine::shrinkLoadForShuffles(Instruction &I) {
   return false;
 }
 
+// Attempt to combine two adjacent fixed-length vector loads, that only feed
+// shufflevector instructions, into a single wider load, rewriting every such
+// shuffle so that operand 0 is the wide load and operand 1 is poison.
+// clang-format off
+// e.g.
+//   %loadA = load <16 x i8>, ptr %a
+//   %gep = getelementptr inbounds <16 x i8>, ptr %a, i64 1
+//   %loadB = load <16 x i8>, ptr %gep
+//   %shuffle0 = shufflevector <16 x i8> %loadA, <16 x i8> %loadB,
+//               <32 x i8> <...>
+//   %shuffle1 = shufflevector <16 x i8> %loadA, <16 x i8> %loadB,
+//               <32 x i8> <...>
+//
+// The fold would transform this to:
+//   %loadAB = load <32 x i8>, ptr %a
+//   %shuffle0 = shufflevector <32 x i8> %loadAB, <32 x i8> poison,
+//               <32 x i8> <...>
+//   %shuffle1 = shufflevector <32 x i8> %loadAB, <32 x i8> poison,
+//               <32 x i8> <...>
+//
+// clang-format on
+// Matching this pattern in codegen becomes difficult and hence, we prefer doing
+// this here.
+bool VectorCombine::foldShuffleOfAdjacentLoads(Instruction &I) {
+  auto *SV = cast<ShuffleVectorInst>(&I);
+
+  // The two operands must be distinct loads of the same fixed vector type.
+  auto *Load0 = dyn_cast<LoadInst>(SV->getOperand(0));
+  auto *Load1 = dyn_cast<LoadInst>(SV->getOperand(1));
+  if (!Load0 || !Load1 || Load0 == Load1 || !Load0->isSimple() ||
+      !Load1->isSimple())
+    return false;
+
+  // Confirm both loads are of fixed vector type.
+  auto *LoadTy = dyn_cast<FixedVectorType>(Load0->getType());
+  if (!LoadTy)
+    return false;
+
+  // We restrict to loads occurring in the same BB for now.
+  if (Load0->getParent() != Load1->getParent())
+    return false;
+
+  if (Load0->getPointerAddressSpace() != Load1->getPointerAddressSpace())
+    return false;
+
+  // Check that the original load type has no padding bits otherwise the wide
+  // load would be incorrect.
+  if (DL->getTypeSizeInBits(LoadTy) != 8 * DL->getTypeStoreSize(LoadTy))
+    return false;
+
+  const int NumElts = LoadTy->getNumElements();
+
+  // Determine which load is at the lower address and confirm the two loads are
+  // exactly contiguous. isConsecutiveAccess(A, B) is true only when B directly
+  // follows A, so we probe both orderings to also handle the reversed case.
+  LoadInst *LowLoad, *HighLoad;
+  assert(SE && "ScalarEvolution is only available for late folds");
+  if (isConsecutiveAccess(Load0, Load1, *DL, *SE)) {
+    LowLoad = Load0;
+    HighLoad = Load1;
+  } else if (isConsecutiveAccess(Load1, Load0, *DL, *SE)) {
+    LowLoad = Load1;
+    HighLoad = Load0;
+  } else {
+    return false;
+  }
+
+  // 1. Check all users of both loads are shuffles.
+  // 2. Check that both loads feed exactly the same set of shuffles.
+  SmallPtrSet<ShuffleVectorInst *, 4> Shuffles;
+  auto AreShufflesOnlyUsersOfLoads = [LowLoad, HighLoad, &Shuffles]() -> bool {
+    // Step 1: collect every user of LowLoad, requiring each to be a shuffle.
+    for (User *U : LowLoad->users()) {
+      auto *SV = dyn_cast<ShuffleVectorInst>(U);
+      if (!SV)
+        return false;
+      Shuffles.insert(SV);
+    }
+
+    // Step 2: every user of HighLoad must be a shuffle already collected from
+    // LowLoad, counting them as we go.
+    unsigned HighLoadUsers = 0;
+    for (User *U : HighLoad->users()) {
+      auto *SV = dyn_cast<ShuffleVectorInst>(U);
+      if (!SV || !Shuffles.contains(SV))
+        return false;
+      ++HighLoadUsers;
+    }
+
+    // Step 3: both loads must feed exactly the same set of shuffles. Combined
+    // with step 2, this guarantees every shuffle uses both LowLoad and
+    // HighLoad, so their operands are exactly {LowLoad, HighLoad}.
+    return HighLoadUsers == Shuffles.size();
+  };
+  if (!AreShufflesOnlyUsersOfLoads())
+    return false;
+
+  // The value loaded by either load must not be clobbered in between the loads.
+  auto *WideTy = FixedVectorType::get(LoadTy->getElementType(), NumElts * 2);
+  LoadInst *FirstLoad = LowLoad, *LastLoad = HighLoad;
+  bool LowComesFirst = LowLoad->comesBefore(HighLoad);
+  if (!LowComesFirst)
+    std::swap(FirstLoad, LastLoad);
+  MemoryLocation FirstLoc = MemoryLocation::get(FirstLoad);
+  if (isMemModifiedBetween(std::next(FirstLoad->getIterator()),
+                           LastLoad->getIterator(), FirstLoc, AA))
+    return false;
+
+  // case 1: wide load = LowLoad + HighLoad   ,
+  //         shuffle 0th operand = LowLoad
+  //         shuffle 1st operand = HighLoad
+  // Implication with this is shuffle mask for the wide load remains unchanged
+  // case 2: wide load = LowLoad + HighLoad   ,
+  //         shuffle 0th operand = HighLoad
+  //         shuffle 1st operand = LowLoad
+  // Implication with this is shuffle mask for the wide load changes
+  auto RemapMask = [LowLoad, HighLoad, NumElts](ShuffleVectorInst *SV,
+                                                SmallVectorImpl<int> &NewMask) {
+    Value *SVOp0 = SV->getOperand(0);
+    Value *SVOp1 = SV->getOperand(1);
+    ArrayRef<int> OldMask = SV->getShuffleMask();
+    assert(((SVOp0 == LowLoad && SVOp1 == HighLoad) ||
+            (SVOp0 == HighLoad && SVOp1 == LowLoad)) &&
+           "Shuffle operands must be exactly {LowLoad, HighLoad} or {HighLoad, "
+           "LowLoad}");
+    unsigned Off0 = SVOp0 == LowLoad ? 0 : NumElts;
+    unsigned Off1 = SVOp1 == LowLoad ? 0 : NumElts;
+
+    NewMask.clear();
+    for (int M : OldMask) {
+      if (M < 0)
+        NewMask.push_back(M);
+      else if (M < NumElts)
+        NewMask.push_back(Off0 + M);
+      else
+        NewMask.push_back(Off1 + (M - NumElts));
+    }
----------------
sushgokh wrote:

done

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