[SLP] Workaround for vectorizing loads with more than one store uses.

[vporpo]

Summary

 L0      L1
/ \     /  \

S0 S0' S1 S1'

In AArch64 this simple pattern of loads with multiple store users is not getting
vectorized because of the cost of the external uses.
This patch is a workaround that reduces the cost of the external use if the def
is a load and the user can be found in the Stores seeds map.

[ABataev]

It is very optimistic estimation. It does not check if the external stores could be vectorized, does not check for the reordering shuffles costs. May lead to regressions in many cases.

[vporpo]

Yeah I agree, it is too optimistic. What if I also check the contents of the tree to make sure it is a store-load only, and also check that the ordering of the external stores matches the loads?

[ABataev]

In D132773#3753642, @vporpo wrote:

Yeah I agree, it is too optimistic. What if I also check the contents of the tree to make sure it is a store-load only, and also check that the ordering of the external stores matches the loads?

Maybe try something like this: for each unsuccessful vectotization attempt record stores, which build vectorized store and compare the uses with this stores. The last stores will be vectorized, because all the users are potentially vectorizable. The process then repeats itself, because of the changes in basicblock. And first stores are also vectorized. You can estimate the reordering cost too in this case.

[davidxl]

Vasileios, Abataev's suggestion seems reasonable. Can you try it?

[vporpo]

@davidxl I tried it, but it is not as simple as it seems. It won't vectorize the code fully unless we also modify the visiting order of the seeds.

[davidxl]

Naive question -- can this be done as a post-processing step with proper reordering?

llvm/include/llvm/Transforms/Vectorize/SLPVectorizer.h

private:
   /// size.
   /// \returns true if a value was vectorized.
   bool tryToVectorizeList(ArrayRef<Value *> VL, slpvectorizer::BoUpSLP &R,
                           bool LimitForRegisterSize = false);
 
   /// Try to vectorize a chain that may start at the operands of \p I.
   bool tryToVectorize(Instruction *I, slpvectorizer::BoUpSLP &R);
 
+  /// \Returns true if \p U is part a store seed.
+  bool isStoreSeed(User *U) const;
+
   /// Vectorize the store instructions collected in Stores.
   bool vectorizeStoreChains(slpvectorizer::BoUpSLP &R);
 
   /// Vectorize the index computations of the getelementptr instructions
   /// collected in GEPs.
   bool vectorizeGEPIndices(BasicBlock *BB, slpvectorizer::BoUpSLP &R);
 
   /// Try to find horizontal reduction or otherwise vectorize a chain of binary

llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp

public:
   Value *vectorizeTree(ExtraValueToDebugLocsMap &ExternallyUsedValues);
 
   /// \returns the cost incurred by unwanted spills and fills, caused by
   /// holding live values over call sites.
   InstructionCost getSpillCost() const;
 
   /// \returns the vectorization cost of the subtree that starts at \p VL.
   /// A negative number means that this is profitable.
-  InstructionCost getTreeCost(ArrayRef<Value *> VectorizedVals = None);
+  InstructionCost getTreeCost(function_ref<bool(User *)> IsSeed,
+                              ArrayRef<Value *> VectorizedVals = None);
 
   /// Construct a vectorizable tree that starts at \p Roots, ignoring users for
   /// the purpose of scheduling and extraction in the \p UserIgnoreLst.
   void buildTree(ArrayRef<Value *> Roots,
                  const SmallDenseSet<Value *> &UserIgnoreLst);
 
   /// Construct a vectorizable tree that starts at \p Roots.
   void buildTree(ArrayRef<Value *> Roots);
for (unsigned I = 0, VF = Mask.size(); I < VF; ++I) {
         Mask[I] = I;
       }
     }
     Prev = Action(Mask, {Prev, Res.first});
   }
   return Prev;
 }
 
-InstructionCost BoUpSLP::getTreeCost(ArrayRef<Value *> VectorizedVals) {
+InstructionCost BoUpSLP::getTreeCost(function_ref<bool(User *)> IsSeed,
+                                     ArrayRef<Value *> VectorizedVals) {
   InstructionCost Cost = 0;
   LLVM_DEBUG(dbgs() << "SLP: Calculating cost for tree of size "
                     << VectorizableTree.size() << ".\n");
 
   unsigned BundleWidth = VectorizableTree[0]->Scalars.size();
 
   for (unsigned I = 0, E = VectorizableTree.size(); I < E; ++I) {
     TreeEntry &TE = *VectorizableTree[I];
for (ExternalUser &EU : ExternalUses) {
     if (isa<FixedVectorType>(EU.Scalar->getType()))
       continue;
 
     // Already counted the cost for external uses when tried to adjust the cost
     // for extractelements, no need to add it again.
     if (isa<ExtractElementInst>(EU.Scalar))
       continue;
 
+    // Skip if this is part of a seed.
+    if (isa<LoadInst>(EU.Scalar) && EU.User && isa<StoreInst>(EU.User) &&
+        IsSeed(EU.User))
+      continue;
+
     // If found user is an insertelement, do not calculate extract cost but try
     // to detect it as a final shuffled/identity match.
     if (auto *VU = dyn_cast_or_null<InsertElementInst>(EU.User)) {
       if (auto *FTy = dyn_cast<FixedVectorType>(VU->getType())) {
         Optional<unsigned> InsertIdx = getInsertIndex(VU);
         if (InsertIdx) {
           const TreeEntry *ScalarTE = getTreeEntry(EU.Scalar);
           auto *It =
bool SLPVectorizerPass::vectorizeStoreChain(ArrayRef<Value *> Chain, BoUpSLP &R,
   if (R.isLoadCombineCandidate())
     return false;
   R.reorderTopToBottom();
   R.reorderBottomToTop();
   R.buildExternalUses();
 
   R.computeMinimumValueSizes();
 
-  InstructionCost Cost = R.getTreeCost();
+  InstructionCost Cost =
+      R.getTreeCost([this](User *U) { return isStoreSeed(U); });
 
   LLVM_DEBUG(dbgs() << "SLP: Found cost = " << Cost << " for VF=" << VF << "\n");
   if (Cost < -SLPCostThreshold) {
     LLVM_DEBUG(dbgs() << "SLP: Decided to vectorize cost = " << Cost << "\n");
 
     using namespace ore;
 
     R.getORE()->emit(OptimizationRemark(SV_NAME, "StoresVectorized",
for (unsigned I = NextInst; I < MaxInst; ++I) {
       R.buildTree(Ops);
       if (R.isTreeTinyAndNotFullyVectorizable())
         continue;
       R.reorderTopToBottom();
       R.reorderBottomToTop(!isa<InsertElementInst>(Ops.front()));
       R.buildExternalUses();
 
       R.computeMinimumValueSizes();
-      InstructionCost Cost = R.getTreeCost();
+      InstructionCost Cost = R.getTreeCost([](User *U) { return false; });
       CandidateFound = true;
       MinCost = std::min(MinCost, Cost);
 
       LLVM_DEBUG(dbgs() << "SLP: Found cost = " << Cost << " for VF=" << VF << "\n");
       if (Cost < -SLPCostThreshold) {
         LLVM_DEBUG(dbgs() << "SLP: Vectorizing list at cost:" << Cost << ".\n");
         R.getORE()->emit(OptimizationRemark(SV_NAME, "VectorizedList",
                                                     cast<Instruction>(Ops[0]))
for (unsigned I = 0, E = ReducedVals.size(); I < E; ++I) {
 
         // Intersect the fast-math-flags from all reduction operations.
         FastMathFlags RdxFMF;
         RdxFMF.set();
         for (Value *U : IgnoreList)
           if (auto *FPMO = dyn_cast<FPMathOperator>(U))
             RdxFMF &= FPMO->getFastMathFlags();
         // Estimate cost.
-        InstructionCost TreeCost = V.getTreeCost(VL);
+        InstructionCost TreeCost =
+            V.getTreeCost([](User *U) { return false; }, VL);
         InstructionCost ReductionCost =
             getReductionCost(TTI, VL, ReduxWidth, RdxFMF);
         InstructionCost Cost = TreeCost + ReductionCost;
         LLVM_DEBUG(dbgs() << "SLP: Found cost = " << Cost << " for reduction\n");
         if (!Cost.isValid()) {
           return nullptr;
         }
         if (Cost >= -SLPCostThreshold) {
for (unsigned BI = 0, BE = Entry.second.size(); BI < BE; BI += MaxElts) {
       // bottom-up phase will be simpler and less costly than building a
       // full-blown top-down phase beginning at the consecutive loads.
       Changed |= tryToVectorizeList(Bundle, R);
     }
   }
   return Changed;
 }
 
+bool SLPVectorizerPass::isStoreSeed(User *U) const {
+  StoreInst *SI = dyn_cast<StoreInst>(U);
+  if (SI == nullptr)
+    return false;
+  Value *Ptr = getUnderlyingObject(SI->getPointerOperand());
+  auto It = Stores.find(Ptr);
+  if (It == Stores.end())
+    return false;
+  const StoreList &SList = It->second;
+  // Check if StoreVL can be found in SList.
+  auto ListIt = find(SList, SI);
+  return ListIt != SList.end();
+}
+
 bool SLPVectorizerPass::vectorizeStoreChains(BoUpSLP &R) {
   bool Changed = false;
   // Sort by type, base pointers and values operand. Value operands must be
   // compatible (have the same opcode, same parent), otherwise it is
   // definitely not profitable to try to vectorize them.
   auto &&StoreSorter = [this](StoreInst *V, StoreInst *V2) {
     if (V->getPointerOperandType()->getTypeID() <
         V2->getPointerOperandType()->getTypeID())

llvm/test/Transforms/SLPVectorizer/AArch64/load-used-by-two-stores.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt -slp-vectorizer -S < %s -mtriple=aarch64-- | FileCheck %s
 
 ; This test check that we can vectorize the simple pattern of a load with more
 ; than one store user:
 ;   L0       L1
 ;  /  \     /  \
 ; S0  S'0  S1  S'1
 
 define void @load_used_by_two_stores_double(double *%ptr) {
 ; CHECK-LABEL: @load_used_by_two_stores_double(
 ; CHECK-NEXT:    [[GEP0:%.*]] = getelementptr inbounds double, double* [[PTR:%.*]], i64 0
-; CHECK-NEXT:    [[GEP1:%.*]] = getelementptr inbounds double, double* [[PTR]], i64 1
-; CHECK-NEXT:    [[LD0:%.*]] = load double, double* [[GEP0]], align 8
-; CHECK-NEXT:    [[LD1:%.*]] = load double, double* [[GEP1]], align 8
-; CHECK-NEXT:    store double [[LD0]], double* [[GEP0]], align 8
-; CHECK-NEXT:    store double [[LD1]], double* [[GEP1]], align 8
-; CHECK-NEXT:    store double [[LD0]], double* [[GEP0]], align 8
-; CHECK-NEXT:    store double [[LD1]], double* [[GEP1]], align 8
+; CHECK-NEXT:    [[TMP1:%.*]] = bitcast double* [[GEP0]] to <2 x double>*
+; CHECK-NEXT:    [[TMP2:%.*]] = load <2 x double>, <2 x double>* [[TMP1]], align 8
+; CHECK-NEXT:    [[TMP3:%.*]] = bitcast double* [[GEP0]] to <2 x double>*
+; CHECK-NEXT:    store <2 x double> [[TMP2]], <2 x double>* [[TMP3]], align 8
+; CHECK-NEXT:    [[TMP4:%.*]] = bitcast double* [[GEP0]] to <2 x double>*
+; CHECK-NEXT:    store <2 x double> [[TMP2]], <2 x double>* [[TMP4]], align 8
 ; CHECK-NEXT:    ret void
 ;
   %gep0 = getelementptr inbounds double, double* %ptr, i64 0
   %gep1 = getelementptr inbounds double, double* %ptr, i64 1
   %ld0 = load double, double *%gep0, align 8
   %ld1 = load double, double *%gep1, align 8
 
   store double %ld0, double *%gep0, align 8
   store double %ld1, double *%gep1, align 8
 
   store double %ld0, double *%gep0, align 8
   store double %ld1, double *%gep1, align 8
   ret void
 }
 
 define void @load_used_by_two_stores_i32(i32 *%ptr) {
 ; CHECK-LABEL: @load_used_by_two_stores_i32(
 ; CHECK-NEXT:    [[GEP0:%.*]] = getelementptr inbounds i32, i32* [[PTR:%.*]], i64 0
-; CHECK-NEXT:    [[GEP1:%.*]] = getelementptr inbounds i32, i32* [[PTR]], i64 1
-; CHECK-NEXT:    [[GEP2:%.*]] = getelementptr inbounds i32, i32* [[PTR]], i64 2
-; CHECK-NEXT:    [[GEP3:%.*]] = getelementptr inbounds i32, i32* [[PTR]], i64 3
-; CHECK-NEXT:    [[LD0:%.*]] = load i32, i32* [[GEP0]], align 8
-; CHECK-NEXT:    [[LD1:%.*]] = load i32, i32* [[GEP1]], align 8
-; CHECK-NEXT:    [[LD2:%.*]] = load i32, i32* [[GEP2]], align 8
-; CHECK-NEXT:    [[LD3:%.*]] = load i32, i32* [[GEP3]], align 8
-; CHECK-NEXT:    store i32 [[LD0]], i32* [[GEP0]], align 8
-; CHECK-NEXT:    store i32 [[LD1]], i32* [[GEP1]], align 8
-; CHECK-NEXT:    store i32 [[LD2]], i32* [[GEP2]], align 8
-; CHECK-NEXT:    store i32 [[LD3]], i32* [[GEP3]], align 8
-; CHECK-NEXT:    store i32 [[LD0]], i32* [[GEP0]], align 8
-; CHECK-NEXT:    store i32 [[LD1]], i32* [[GEP1]], align 8
-; CHECK-NEXT:    store i32 [[LD2]], i32* [[GEP2]], align 8
-; CHECK-NEXT:    store i32 [[LD3]], i32* [[GEP3]], align 8
+; CHECK-NEXT:    [[TMP1:%.*]] = bitcast i32* [[GEP0]] to <4 x i32>*
+; CHECK-NEXT:    [[TMP2:%.*]] = load <4 x i32>, <4 x i32>* [[TMP1]], align 8
+; CHECK-NEXT:    [[TMP3:%.*]] = bitcast i32* [[GEP0]] to <4 x i32>*
+; CHECK-NEXT:    store <4 x i32> [[TMP2]], <4 x i32>* [[TMP3]], align 8
+; CHECK-NEXT:    [[TMP4:%.*]] = bitcast i32* [[GEP0]] to <4 x i32>*
+; CHECK-NEXT:    store <4 x i32> [[TMP2]], <4 x i32>* [[TMP4]], align 8
 ; CHECK-NEXT:    ret void
 ;
   %gep0 = getelementptr inbounds i32, i32* %ptr, i64 0
   %gep1 = getelementptr inbounds i32, i32* %ptr, i64 1
   %gep2 = getelementptr inbounds i32, i32* %ptr, i64 2
   %gep3 = getelementptr inbounds i32, i32* %ptr, i64 3
 
   %ld0 = load i32, i32 *%gep0, align 8
;
   store i32 %ld3, i32 *%gep3, align 8
 
   ret void
 }
 
 define void @load_used_by_two_stores_i8(i8 *%ptr) {
 ; CHECK-LABEL: @load_used_by_two_stores_i8(
 ; CHECK-NEXT:    [[GEP0:%.*]] = getelementptr inbounds i8, i8* [[PTR:%.*]], i64 0
-; CHECK-NEXT:    [[GEP1:%.*]] = getelementptr inbounds i8, i8* [[PTR]], i64 1
-; CHECK-NEXT:    [[GEP2:%.*]] = getelementptr inbounds i8, i8* [[PTR]], i64 2
-; CHECK-NEXT:    [[GEP3:%.*]] = getelementptr inbounds i8, i8* [[PTR]], i64 3
-; CHECK-NEXT:    [[GEP4:%.*]] = getelementptr inbounds i8, i8* [[PTR]], i64 4
-; CHECK-NEXT:    [[GEP5:%.*]] = getelementptr inbounds i8, i8* [[PTR]], i64 5
-; CHECK-NEXT:    [[GEP6:%.*]] = getelementptr inbounds i8, i8* [[PTR]], i64 6
-; CHECK-NEXT:    [[GEP7:%.*]] = getelementptr inbounds i8, i8* [[PTR]], i64 7
-; CHECK-NEXT:    [[GEP8:%.*]] = getelementptr inbounds i8, i8* [[PTR]], i64 8
-; CHECK-NEXT:    [[GEP9:%.*]] = getelementptr inbounds i8, i8* [[PTR]], i64 9
-; CHECK-NEXT:    [[GEP10:%.*]] = getelementptr inbounds i8, i8* [[PTR]], i64 10
-; CHECK-NEXT:    [[GEP11:%.*]] = getelementptr inbounds i8, i8* [[PTR]], i64 11
-; CHECK-NEXT:    [[GEP12:%.*]] = getelementptr inbounds i8, i8* [[PTR]], i64 12
-; CHECK-NEXT:    [[GEP13:%.*]] = getelementptr inbounds i8, i8* [[PTR]], i64 13
-; CHECK-NEXT:    [[GEP14:%.*]] = getelementptr inbounds i8, i8* [[PTR]], i64 14
-; CHECK-NEXT:    [[GEP15:%.*]] = getelementptr inbounds i8, i8* [[PTR]], i64 15
-; CHECK-NEXT:    [[LD0:%.*]] = load i8, i8* [[GEP0]], align 8
-; CHECK-NEXT:    [[LD1:%.*]] = load i8, i8* [[GEP1]], align 8
-; CHECK-NEXT:    [[LD2:%.*]] = load i8, i8* [[GEP2]], align 8
-; CHECK-NEXT:    [[LD3:%.*]] = load i8, i8* [[GEP3]], align 8
-; CHECK-NEXT:    [[LD4:%.*]] = load i8, i8* [[GEP4]], align 8
-; CHECK-NEXT:    [[LD5:%.*]] = load i8, i8* [[GEP5]], align 8
-; CHECK-NEXT:    [[LD6:%.*]] = load i8, i8* [[GEP6]], align 8
-; CHECK-NEXT:    [[LD7:%.*]] = load i8, i8* [[GEP7]], align 8
-; CHECK-NEXT:    [[LD8:%.*]] = load i8, i8* [[GEP8]], align 8
-; CHECK-NEXT:    [[LD9:%.*]] = load i8, i8* [[GEP9]], align 8
-; CHECK-NEXT:    [[LD10:%.*]] = load i8, i8* [[GEP10]], align 8
-; CHECK-NEXT:    [[LD11:%.*]] = load i8, i8* [[GEP11]], align 8
-; CHECK-NEXT:    [[LD12:%.*]] = load i8, i8* [[GEP12]], align 8
-; CHECK-NEXT:    [[LD13:%.*]] = load i8, i8* [[GEP13]], align 8
-; CHECK-NEXT:    [[LD14:%.*]] = load i8, i8* [[GEP14]], align 8
-; CHECK-NEXT:    [[LD15:%.*]] = load i8, i8* [[GEP15]], align 8
-; CHECK-NEXT:    store i8 [[LD0]], i8* [[GEP0]], align 8
-; CHECK-NEXT:    store i8 [[LD1]], i8* [[GEP1]], align 8
-; CHECK-NEXT:    store i8 [[LD2]], i8* [[GEP2]], align 8
-; CHECK-NEXT:    store i8 [[LD3]], i8* [[GEP3]], align 8
-; CHECK-NEXT:    store i8 [[LD4]], i8* [[GEP4]], align 8
-; CHECK-NEXT:    store i8 [[LD5]], i8* [[GEP5]], align 8
-; CHECK-NEXT:    store i8 [[LD6]], i8* [[GEP6]], align 8
-; CHECK-NEXT:    store i8 [[LD7]], i8* [[GEP7]], align 8
-; CHECK-NEXT:    store i8 [[LD8]], i8* [[GEP8]], align 8
-; CHECK-NEXT:    store i8 [[LD9]], i8* [[GEP9]], align 8
-; CHECK-NEXT:    store i8 [[LD10]], i8* [[GEP10]], align 8
-; CHECK-NEXT:    store i8 [[LD11]], i8* [[GEP11]], align 8
-; CHECK-NEXT:    store i8 [[LD12]], i8* [[GEP12]], align 8
-; CHECK-NEXT:    store i8 [[LD13]], i8* [[GEP13]], align 8
-; CHECK-NEXT:    store i8 [[LD14]], i8* [[GEP14]], align 8
-; CHECK-NEXT:    store i8 [[LD15]], i8* [[GEP15]], align 8
-; CHECK-NEXT:    store i8 [[LD0]], i8* [[GEP0]], align 8
-; CHECK-NEXT:    store i8 [[LD1]], i8* [[GEP1]], align 8
-; CHECK-NEXT:    store i8 [[LD2]], i8* [[GEP2]], align 8
-; CHECK-NEXT:    store i8 [[LD3]], i8* [[GEP3]], align 8
-; CHECK-NEXT:    store i8 [[LD4]], i8* [[GEP4]], align 8
-; CHECK-NEXT:    store i8 [[LD5]], i8* [[GEP5]], align 8
-; CHECK-NEXT:    store i8 [[LD6]], i8* [[GEP6]], align 8
-; CHECK-NEXT:    store i8 [[LD7]], i8* [[GEP7]], align 8
-; CHECK-NEXT:    store i8 [[LD8]], i8* [[GEP8]], align 8
-; CHECK-NEXT:    store i8 [[LD9]], i8* [[GEP9]], align 8
-; CHECK-NEXT:    store i8 [[LD10]], i8* [[GEP10]], align 8
-; CHECK-NEXT:    store i8 [[LD11]], i8* [[GEP11]], align 8
-; CHECK-NEXT:    store i8 [[LD12]], i8* [[GEP12]], align 8
-; CHECK-NEXT:    store i8 [[LD13]], i8* [[GEP13]], align 8
-; CHECK-NEXT:    store i8 [[LD14]], i8* [[GEP14]], align 8
-; CHECK-NEXT:    store i8 [[LD15]], i8* [[GEP15]], align 8
+; CHECK-NEXT:    [[TMP1:%.*]] = bitcast i8* [[GEP0]] to <16 x i8>*
+; CHECK-NEXT:    [[TMP2:%.*]] = load <16 x i8>, <16 x i8>* [[TMP1]], align 8
+; CHECK-NEXT:    [[TMP3:%.*]] = bitcast i8* [[GEP0]] to <16 x i8>*
+; CHECK-NEXT:    store <16 x i8> [[TMP2]], <16 x i8>* [[TMP3]], align 8
+; CHECK-NEXT:    [[TMP4:%.*]] = bitcast i8* [[GEP0]] to <16 x i8>*
+; CHECK-NEXT:    store <16 x i8> [[TMP2]], <16 x i8>* [[TMP4]], align 8
 ; CHECK-NEXT:    ret void
 ;
   %gep0 = getelementptr inbounds i8, i8* %ptr, i64 0
   %gep1 = getelementptr inbounds i8, i8* %ptr, i64 1
   %gep2 = getelementptr inbounds i8, i8* %ptr, i64 2
   %gep3 = getelementptr inbounds i8, i8* %ptr, i64 3
   %gep4 = getelementptr inbounds i8, i8* %ptr, i64 4
   %gep5 = getelementptr inbounds i8, i8* %ptr, i64 5