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

[Ryan Buchner 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;
+
+    SmallDenseSet<Value *> UserIgnoreList;
+    for (Instruction *SI : StoresToReplace)
+      UserIgnoreList.insert(SI);
+    for (Instruction *I : BuildVectorInsts)
+      UserIgnoreList.insert(I);
+
+    SmallVector<SmallVector<Value *, 4>, 4> Groups;
+    SmallVector<int64_t, 4> GroupOffsets;
+    for (int64_t Start = MinOffset; Start <= MaxOffset; Start += VF) {
+      Groups.emplace_back();
+      GroupOffsets.push_back(Start);
+      for (int64_t Offset = Start; Offset < Start + static_cast<int64_t>(VF);
+           ++Offset)
+        Groups.back().push_back(Lanes[Offset].first);
+    }
+
+    bool AllProfitable = true;
+    for (ArrayRef<Value *> Group : Groups) {
+      R.buildTree(Group, UserIgnoreList);
+      if (R.isTreeTinyAndNotFullyVectorizable()) {
+        AllProfitable = false;
+        break;
+      }
+      if (R.isProfitableToReorder()) {
+        R.reorderTopToBottom();
+        R.reorderBottomToTop();
+      }
+      R.transformNodes();
+      R.computeMinimumValueSizes();
+      InstructionCost TreeCost = R.calculateTreeCostAndTrimNonProfitable();
+      R.buildExternalUses();
+      InstructionCost Cost = R.getTreeCost(TreeCost);
+      if (Cost >= -SLPCostThreshold) {
+        AllProfitable = false;
+        break;
+      }
+    }
+    if (!AllProfitable)
+      continue;
+
+    IRBuilder<> Builder(VecStore);
+    Type *I8Ty = Type::getInt8Ty(VecStore->getContext());
+    for (auto [Idx, Group] : enumerate(Groups)) {
+      R.buildTree(Group, UserIgnoreList);
----------------
bababuck wrote:

If I'm understanding correctly, we build all of the trees and check their cost. If we decide that we're going to vectorize all of them, we then come back and build the trees again and actually vectorize them.

Is there any reason this second time we build the trees we don't reorder/transform/trim/etc?

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