[llvm] [SLP] Split blocking build-vector stores in scalar chains (PR #194970)

[Alexey Bataev via llvm-commits]

================
@@ -27265,6 +27267,183 @@ bool SLPVectorizerPass::vectorizeStores(
   return Changed;
 }
 
+static bool collectBuildVector(Value *V, SmallVectorImpl<Value *> &Elts,
+                               SmallVectorImpl<Instruction *> *Insts = nullptr) {
+  auto *VecTy = dyn_cast<FixedVectorType>(V->getType());
+  if (!VecTy)
+    return false;
+  Elts.assign(VecTy->getNumElements(), nullptr);
+  Value *Cur = V;
+  while (auto *IE = dyn_cast<InsertElementInst>(Cur)) {
+    auto *Idx = dyn_cast<ConstantInt>(IE->getOperand(2));
+    if (!Idx || Idx->getValue().uge(Elts.size()))
+      return false;
+    Elts[Idx->getZExtValue()] = IE->getOperand(1);
+    if (Insts)
+      Insts->push_back(IE);
+    Cur = IE->getOperand(0);
+  }
+  if (!isa<PoisonValue, UndefValue>(Cur))
+    return false;
+  return all_of(Elts, [](Value *Elt) { return Elt != nullptr; });
+}
+
+/// Retile a scalar-store / build-vector-store / scalar-store sequence without
+/// first materializing scalar stores in the IR. The build-vector store is
+/// treated as covering its lane range, and BoUpSLP decides whether each final
+/// full-width store group is profitable.
+bool SLPVectorizerPass::vectorizeBuildVectorStoreChains(BasicBlock *BB,
+                                                        BoUpSLP &R) {
+  bool Changed = false;
+  for (Instruction &I : make_early_inc_range(*BB)) {
+    auto *VecStore = dyn_cast<StoreInst>(&I);
+    if (!VecStore || !VecStore->isSimple() || R.isDeleted(VecStore))
+      continue;
+
+    SmallVector<Value *, 16> VecElts;
+    SmallVector<Instruction *, 16> BuildVectorInsts;
+    if (!collectBuildVector(VecStore->getValueOperand(), VecElts,
+                            &BuildVectorInsts))
+      continue;
+
+    Type *EltTy = VecElts.front()->getType();
+    TypeSize EltSize = DL->getTypeStoreSize(EltTy);
+    if (EltSize.isScalable())
+      continue;
+
+    unsigned VF = VecElts.size();
+    Value *UnderlyingObject = getUnderlyingObject(VecStore->getPointerOperand());
+    MapVector<int64_t, std::pair<Value *, StoreInst *>> Lanes;
+    for (auto [Idx, V] : enumerate(VecElts))
+      Lanes[Idx] = {V, VecStore};
+
+    for (Instruction &J : *BB) {
+      auto *SI = dyn_cast<StoreInst>(&J);
+      if (!SI || SI == VecStore || !SI->isSimple() || R.isDeleted(SI) ||
+          SI->getValueOperand()->getType() != EltTy ||
+          getUnderlyingObject(SI->getPointerOperand()) != UnderlyingObject)
+        continue;
+
+      std::optional<int64_t> Diff =
+          getPointersDiff(EltTy, VecStore->getPointerOperand(), EltTy,
+                          SI->getPointerOperand(), *DL, *SE,
+                          /*StrictCheck=*/true);
+      if (!Diff || (*Diff >= 0 && *Diff < static_cast<int64_t>(VF)))
+        continue;
+      Lanes[*Diff] = {SI->getValueOperand(), SI};
+    }
+
+    int64_t MinOffset = Lanes.front().first;
+    int64_t MaxOffset = Lanes.front().first;
+    for (auto &Lane : Lanes) {
+      MinOffset = std::min(MinOffset, Lane.first);
+      MaxOffset = std::max(MaxOffset, Lane.first);
+    }
+    if (Lanes.size() <= VF || MinOffset >= 0 ||
+        MaxOffset < static_cast<int64_t>(VF))
+      continue;
+    if (static_cast<unsigned>(MaxOffset - MinOffset + 1) != Lanes.size() ||
+        Lanes.size() % VF != 0)
+      continue;
+
+    SmallVector<Instruction *> StoresToReplace;
+    for (auto &Lane : Lanes)
+      StoresToReplace.push_back(Lane.second.second);
+    sort(StoresToReplace);
+    StoresToReplace.erase(llvm::unique(StoresToReplace), StoresToReplace.end());
+
+    Instruction *FirstStore = StoresToReplace.front();
+    Instruction *LastStore = StoresToReplace.front();
+    for (Instruction *SI : StoresToReplace) {
+      if (SI->comesBefore(FirstStore))
+        FirstStore = SI;
+      if (LastStore->comesBefore(SI))
+        LastStore = SI;
+    }
+
+    bool SawFirst = false;
+    bool SafeToMove = true;
+    for (Instruction &J : *BB) {
+      if (&J == FirstStore)
+        SawFirst = true;
+      if (!SawFirst)
+        continue;
+      if (J.mayReadOrWriteMemory() && !is_contained(StoresToReplace, &J)) {
+        SafeToMove = false;
+        break;
+      }
+      if (&J == LastStore)
+        break;
+    }
+    if (!SafeToMove)
+      continue;
----------------
alexey-bataev wrote:

It is against SLP design. It does not perform any transformations before the modelling/estimation.

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