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

[Alexey Bataev via llvm-commits]

================
@@ -27302,20 +27240,137 @@ bool SLPVectorizerPass::vectorizeStores(
                         return A && (!B || A->getStride() < B->getStride());
                       });
 
+    auto TryVectorizeVectorLaneContext = [&](StoreChainContext &Context,
+                                             unsigned VF) {
+      ArrayRef<StoreLane> Lanes = Context.getOperands();
+      bool HasScalarLane = any_of(
+          Lanes, [](const StoreLane &Lane) { return !Lane.IsVectorLane; });
+      // Keep the initial support transactional: every lane must belong to a
+      // full output group, so no original vector store is partly replaced.
+      if (!Context.hasVectorLane() || !HasScalarLane || Lanes.size() % VF != 0)
+        return false;
+
+      // A vector store can cover lanes that land in several output groups. For
+      // example, p[3..6] may be replaced by new stores p[0..3] and p[4..7].
+      // Only erase the original stores after every affected group is
+      // profitable.
+      SmallVector<Instruction *> StoresToReplace;
+      SmallDenseSet<Value *> UserIgnoreList;
+      for (const StoreLane &Lane : Lanes) {
+        if (is_contained(StoresToReplace, Lane.Store))
+          continue;
+        StoresToReplace.push_back(Lane.Store);
+        UserIgnoreList.insert(Lane.Store);
+        if (!Lane.IsVectorLane)
+          continue;
+        SmallVector<Value *, 16> Elts;
+        SmallVector<Instruction *, 16> BuildVectorInsts;
+        if (!collectBuildVector(Lane.Store->getValueOperand(), Elts,
+                                BuildVectorInsts))
+          return false;
+        for (Instruction *I : BuildVectorInsts)
+          UserIgnoreList.insert(I);
+      }
+
+      TypeSize EltSize =
+          DL->getTypeStoreSize(Lanes.front().ScalarValue->getType());
+      if (EltSize.isScalable())
+        return false;
+      Type *I8Ty = Type::getInt8Ty(Lanes.front().Store->getContext());
+
+      auto TryBuildStoreValue = [&](ArrayRef<StoreLane> Group,
+                                    Value **VectorizedValue = nullptr) {
+        SmallVector<Value *, 4> Values;
+        for (const StoreLane &Lane : Group)
+          Values.push_back(Lane.ScalarValue);
+        R.buildTree(Values, UserIgnoreList, /*IsReduction=*/false);
+        if (R.isTreeTinyAndNotFullyVectorizable())
+          return false;
+        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)
+          return false;
+        if (VectorizedValue)
+          *VectorizedValue = R.vectorizeTree();
+        return true;
+      };
+
+      for (unsigned I = 0, E = Lanes.size(); I != E; I += VF)
+        if (!TryBuildStoreValue(Lanes.slice(I, VF)))
+          return false;
+
+      for (unsigned I = 0, E = Lanes.size(); I != E; I += VF) {
+        ArrayRef<StoreLane> Group = Lanes.slice(I, VF);
+        Value *Vec = nullptr;
+        bool Profitable = TryBuildStoreValue(Group, &Vec);
+        assert(Profitable && Vec && "Expected profitable group from dry run");
+        (void)Profitable;
+
+        IRBuilder<> Builder(Group.front().Store);
+        Value *Ptr = Builder.CreateConstGEP1_64(
+            I8Ty, Group.front().Store->getPointerOperand(),
+            Group.front().Lane * EltSize.getFixedValue());
+        Align StoreAlign = Group.front().getAlignment(EltSize);
+        int64_t GroupOffset = Group.front().Offset;
+        for (const StoreLane &Lane : Group.drop_front()) {
+          Align LaneAlignAtGroupBase = commonAlignment(
+              Lane.getAlignment(EltSize),
+              (Lane.Offset - GroupOffset) * EltSize.getFixedValue());
+          StoreAlign = std::min(StoreAlign, LaneAlignAtGroupBase);
+        }
+        Builder.CreateAlignedStore(Vec, Ptr, StoreAlign);
+      }
+
+      for (Instruction *SI : StoresToReplace)
+        R.eraseInstruction(SI);
+      for (Value *V : UserIgnoreList)
+        if (auto *I = dyn_cast<Instruction>(V);
+            I && !R.isDeleted(I) && isInstructionTriviallyDead(I))
+          R.eraseInstruction(I);
+      VectorizedStores.insert_range(StoresToReplace);
+      Changed = true;
+      return true;
+    };
+
     for (unsigned LimitVF = GlobalMaxVF; LimitVF > 0;
          LimitVF = bit_ceil(LimitVF) / 2) {
       for (auto &CtxPtr : AllContexts) {
         if (!CtxPtr)
           break;
         StoreChainContext &Context = *CtxPtr;
+        if (Context.hasVectorLane()) {
+          if (std::optional<unsigned> VF = Context.getCurrentVF();
+              VF && *VF >= LimitVF) {
----------------
alexey-bataev wrote:

```suggestion
              VF.value_or(0) >= LimitVF) {
```

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