[SLP] Make SLPVectorizer to use `llvm.masked.scatter` intrinsic

[anton-afanasyev]

Summary

This is symmetrical to fcad8d3635c commit adding llvm.masked.gather.
For the scattered operands of store instructions it makes sense
to use scattered store intrinsic, which can lower to native instruction
for X86/AVX512 and ARM/SVE. This also enables building
vectorization tree with entries containing scattered operands.

Actually NFC now, since stores could be only seed entry, but we
collect only consecutive stores there for now. It will be changed
in future commits.

[anton-afanasyev]

Analogue of https://reviews.llvm.org/D90445 for stores.

[anton-afanasyev]

Align fix. Style fix.

[RKSimon]

Actually NFC now, since stores could be only seed entry, but we collect only consecutive stores there for now. It will be changed in future commits.

Are you looking at supporting non-consecutive stores? This patch probably shouldn't be reviewed until we actually have something that uses and tests it.

[anton-afanasyev]

In D91919#2439339, @RKSimon wrote:

Actually NFC now, since stores could be only seed entry, but we collect only consecutive stores there for now. It will be changed in future commits.

Are you looking at supporting non-consecutive stores? This patch probably shouldn't be reviewed until we actually have something that uses and tests it.

Yes, I've planned to add such stores generation at the next step. But you're right, better to do this together with this step.

[xbolva00]

Please add tests

llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp

case Instruction::Load: {
         TreeEntry *TE = newTreeEntry(VL, TreeEntry::ScatterVectorize, Bundle, S,
                                      UserTreeIdx, ReuseShuffleIndicies);
         TE->setOperandsInOrder();
         buildTree_rec(PointerOps, Depth + 1, {TE, 0});
         LLVM_DEBUG(dbgs() << "SLP: added a vector of non-consecutive loads.\n");
         return;
       }
 
-      LLVM_DEBUG(dbgs() << "SLP: Gathering non-consecutive loads.\n");
+      LLVM_DEBUG(dbgs() << "SLP: Gathering non-vectorized loads.\n");
       BS.cancelScheduling(VL, VL0);
       newTreeEntry(VL, None /*not vectorized*/, S, UserTreeIdx,
                    ReuseShuffleIndicies);
       return;
     }
     case Instruction::ZExt:
     case Instruction::SExt:
     case Instruction::FPToUI:
case Instruction::Store: {
             TE->setOperandsInOrder();
             buildTree_rec(Operands, Depth + 1, {TE, 0});
             LLVM_DEBUG(dbgs() << "SLP: added a vector of jumbled stores.\n");
             findRootOrder(CurrentOrder);
             ++NumOpsWantToKeepOrder[CurrentOrder];
           }
           return;
         }
+        // Vectorizing non-consecutive stores with `llvm.masked.scatter`.
+        TreeEntry *TE = newTreeEntry(VL, TreeEntry::ScatterVectorize, Bundle, S,
+                                     UserTreeIdx, ReuseShuffleIndicies);
+        TE->setOperandsInOrder();
+        buildTree_rec(PointerOps, Depth + 1, {TE, 0});
+        LLVM_DEBUG(
+            dbgs() << "SLP: added a vector of non-consecutive stores.\n");
+        return;
       }
 
       BS.cancelScheduling(VL, VL0);
       newTreeEntry(VL, None /*not vectorized*/, S, UserTreeIdx,
                    ReuseShuffleIndicies);
-      LLVM_DEBUG(dbgs() << "SLP: Non-consecutive store.\n");
+      LLVM_DEBUG(dbgs() << "SLP: Gathering non-vectorized store.\n");
       return;
     }
     case Instruction::Call: {
       // Check if the calls are all to the same vectorizable intrinsic or
       // library function.
       CallInst *CI = cast<CallInst>(VL0);
       Intrinsic::ID ID = getVectorIntrinsicIDForCall(CI, TLI);
 
case Instruction::GetElementPtr: {
       int ScalarCost = VecTy->getNumElements() * ScalarEltCost;
       int VecCost =
           TTI->getArithmeticInstrCost(Instruction::Add, VecTy, CostKind,
                                       Op1VK, Op2VK);
       return ReuseShuffleCost + VecCost - ScalarCost;
     }
     case Instruction::Load: {
       // Cost of wide load - cost of scalar loads.
-      Align alignment = cast<LoadInst>(VL0)->getAlign();
-      int ScalarEltCost =
-          TTI->getMemoryOpCost(Instruction::Load, ScalarTy, alignment, 0,
-                               CostKind, VL0);
+      Align Alignment = cast<LoadInst>(VL0)->getAlign();
+      int ScalarEltCost = TTI->getMemoryOpCost(Instruction::Load, ScalarTy,
+                                               Alignment, 0, CostKind, VL0);
       if (NeedToShuffleReuses) {
         ReuseShuffleCost -= (ReuseShuffleNumbers - VL.size()) * ScalarEltCost;
       }
       int ScalarLdCost = VecTy->getNumElements() * ScalarEltCost;
       int VecLdCost;
       if (E->State == TreeEntry::Vectorize) {
-        VecLdCost = TTI->getMemoryOpCost(Instruction::Load, VecTy, alignment, 0,
+        VecLdCost = TTI->getMemoryOpCost(Instruction::Load, VecTy, Alignment, 0,
                                          CostKind, VL0);
       } else {
         assert(E->State == TreeEntry::ScatterVectorize && "Unknown EntryState");
+        Align CommonAlignment = Alignment;
+        for (Value *V : VL)
+          CommonAlignment =
+              commonAlignment(CommonAlignment, cast<LoadInst>(V)->getAlign());
         VecLdCost = TTI->getGatherScatterOpCost(
             Instruction::Load, VecTy, cast<LoadInst>(VL0)->getPointerOperand(),
-            /*VariableMask=*/false, alignment, CostKind, VL0);
+            /*VariableMask=*/false, CommonAlignment, CostKind, VL0);
       }
       if (!E->ReorderIndices.empty()) {
         // TODO: Merge this shuffle with the ReuseShuffleCost.
         VecLdCost += TTI->getShuffleCost(
             TargetTransformInfo::SK_PermuteSingleSrc, VecTy);
       }
       return ReuseShuffleCost + VecLdCost - ScalarLdCost;
     }
     case Instruction::Store: {
       // We know that we can merge the stores. Calculate the cost.
       bool IsReorder = !E->ReorderIndices.empty();
       auto *SI =
           cast<StoreInst>(IsReorder ? VL[E->ReorderIndices.front()] : VL0);
       Align Alignment = SI->getAlign();
       int ScalarEltCost =
           TTI->getMemoryOpCost(Instruction::Store, ScalarTy, Alignment, 0,
                                CostKind, VL0);
       if (NeedToShuffleReuses)
         ReuseShuffleCost = -(ReuseShuffleNumbers - VL.size()) * ScalarEltCost;
       int ScalarStCost = VecTy->getNumElements() * ScalarEltCost;
-      int VecStCost = TTI->getMemoryOpCost(Instruction::Store,
-                                           VecTy, Alignment, 0, CostKind, VL0);
+      int VecStCost;
+      if (E->State == TreeEntry::Vectorize) {
+        VecStCost = TTI->getMemoryOpCost(Instruction::Store, VecTy, Alignment,
+                                         0, CostKind, VL0);
+      } else {
+        assert(E->State == TreeEntry::ScatterVectorize && "Unhandled state");
+        Align CommonAlignment = Alignment;
+        for (Value *V : VL)
+          CommonAlignment =
+              commonAlignment(CommonAlignment, cast<StoreInst>(V)->getAlign());
+        VecStCost = TTI->getGatherScatterOpCost(
+            Instruction::Store, VecTy,
+            cast<StoreInst>(VL0)->getPointerOperand(),
+            /*VariableMask=*/false, CommonAlignment, CostKind, VL0);
+      }
       if (IsReorder) {
         // TODO: Merge this shuffle with the ReuseShuffleCost.
         VecStCost += TTI->getShuffleCost(
             TargetTransformInfo::SK_PermuteSingleSrc, VecTy);
       }
       return ReuseShuffleCost + VecStCost - ScalarStCost;
     }
     case Instruction::Call: {
switch (ShuffleOrOp) {
     case Instruction::Store: {
       bool IsReorder = !E->ReorderIndices.empty();
       auto *SI = cast<StoreInst>(
           IsReorder ? E->Scalars[E->ReorderIndices.front()] : VL0);
       unsigned AS = SI->getPointerAddressSpace();
 
       setInsertPointAfterBundle(E);
 
+      Instruction *NewSI;
       Value *VecValue = vectorizeTree(E->getOperand(0));
       if (IsReorder) {
         SmallVector<int, 4> Mask(E->ReorderIndices.begin(),
                                  E->ReorderIndices.end());
         VecValue = Builder.CreateShuffleVector(VecValue, Mask, "reorder_shuf");
       }
+      if (E->State == TreeEntry::Vectorize) {
         Value *ScalarPtr = SI->getPointerOperand();
         Value *VecPtr = Builder.CreateBitCast(
             ScalarPtr, VecValue->getType()->getPointerTo(AS));
-      StoreInst *ST = Builder.CreateAlignedStore(VecValue, VecPtr,
-                                                 SI->getAlign());
+        NewSI = Builder.CreateAlignedStore(VecValue, VecPtr, SI->getAlign());
 
         // The pointer operand uses an in-tree scalar, so add the new BitCast to
         // ExternalUses to make sure that an extract will be generated in the
         // future.
         if (getTreeEntry(ScalarPtr))
-        ExternalUses.push_back(ExternalUser(ScalarPtr, cast<User>(VecPtr), 0));
+          ExternalUses.push_back(
+              ExternalUser(ScalarPtr, cast<User>(VecPtr), 0));
+      } else {
+        assert(E->State == TreeEntry::ScatterVectorize && "Unhandled state");
+        Value *VecPtr = vectorizeTree(E->getOperand(1));
+        // Use the minimum alignment of the scattered stores.
+        Align CommonAlignment = SI->getAlign();
+        for (Value *V : E->Scalars)
+          CommonAlignment =
+              commonAlignment(CommonAlignment, cast<StoreInst>(V)->getAlign());
+        NewSI = Builder.CreateMaskedScatter(VecValue, VecPtr, CommonAlignment);
+      }
 
-      Value *V = propagateMetadata(ST, E->Scalars);
+      Value *V = propagateMetadata(NewSI, E->Scalars);
       if (NeedToShuffleReuses)
         V = Builder.CreateShuffleVector(V, E->ReuseShuffleIndices, "shuffle");
 
       E->VectorizedValue = V;
       ++NumVectorInstructions;
       return V;
     }
     case Instruction::GetElementPtr: {