[SLP] Vectorize gather-like idioms ending at non-consecutive loads.

[mssimpso]

Summary

This patch tries to vectorize gather-like expression trees ending at
non-consecutive loads, such as the one shown in the example below.

... = g[a[0] - b[0]] + g[a[1] - b[1]] + ... + g[a[n] - b[n]];

Here, the index calculations for the "g" accesses can be vectorized. The loads
of the "a" and "b" array elements and the subtractions can all be replaced by
their vector equivalents. Our bottom-up vectorizer currently misses cases like
this because the expression trees don't end in stores or reductions.

It's possible to vectorize these cases in a top-down phase beginning at the
consecutive loads. However, I've chosen here to detect the specific pattern of
interest and proceed bottom-up as we do with other interesting cases. The
advantage of this approach is that it avoids the complexity, compile-time, and
phase ordering issues of a full-blown top-down pass. The disadvantage is that
it's probably not as general as it would be otherwise.

The primary changes included in the patch allow us to (1) vectorize the
gather-like pattern shown above and (2) set vector factors based on the width
of memory accesses in the expression trees. Your feedback is welcome.

[nadav]

Matthew, thanks for working on this.

In your patch you are vectorizing very short trees that are made out of loads, geps and potentially stores. In the early days of the SLP we disabled the vectorization of short trees because of two reasons:

1. It is very difficult to estimate the profitability of this kind of vectorization. The trees are so small and the scores of the trees are very low, often close to zero. Estimating the cost of gathers is even more difficult. I don't know if you’ve ran performance measurements but I’d be surprised if this kind of vectorization would be profitable for general programs.
2. The SelectionDAG ConsecutiveStore optimization already vectorizes short load-store trees.

I think that it is better to optimize this pattern by extending SelectionDAG’s ConsecutiveStore optimizations.

What do you think?

Thanks,
Nadav

[mssimpso]

Hi Nadav,

Thanks very much for the quick feedback! I'm happy to consider a different implementation. I have have some questions for you first, though, if you don't mind.

I don't think the expression trees this patch affects are necessarily limited in size. The requirement is that they be seeded by GEPs (that are used by non-consecutive loads). As far as I know, the trees can be arbitrarily large. The trees in my example are small (gep, zext, sub, load), but this doesn't have to be the case in general. Are you suggesting that in your experience, seeding SLP with these GEPs typically only hits short, unprofitable trees in practice?

I have measured performance on our workloads. The cost estimate for the test case I provided in the patch comes in well below zero and is profitable. Estimating the cost of gathers is difficult, but here we are only optimizing the index calculations. I will admit that the number of programs we care about is limited. It would nice to have some additional data points. Is this something you could help with?

Finally, I'm not familiar with the portion of SelectionDAG you mention, so please excuse the rest of my questions. First, there aren't any stores in my example, so would the ConsecutiveStores optimizations you mention even apply? And lastly, wouldn't we still be subject to the same cost estimates (and imprecision therein) if the implementation was moved elsewhere?

Thanks again!

[hfinkel]

I don't think the expression trees this patch affects are necessarily limited in size.

I agree; these patterns can occur with reasonable complexity in practice.

Have you run this on the test suite, did you find any statistically-significant compile-time impact?

lib/Transforms/Vectorize/SLPVectorizer.cpp

4294

prpofitable (typo)

[mssimpso]

Nadav/Hal,

Thanks very much for the additional feedback. Just letting you know that I haven't forgotten about this review. I will provide some compile-time and performance results soon.

[mssimpso]

Fixed typo in comment.

[mssimpso]

Nadav/Hal,

Below are the statistically significant compile-time differences observed for the test suite, spec2000, and spec2006 (there is one). Results were computed from 10 samples each, median aggregation, 95% confidence intervals, and a 0.05 statistical significance level for the Mann–Whitney U test.

Program                                                             Base  Change      %
---------------------------------------------------------------------------------------
MultiSource/Benchmarks/MiBench/security-rijndael/security-rijndael  0.73    0.89  18.88

No statistically significant performance differences were observed for spec2000 and spec2006 on a Cortex-A57-like cpu (but see the explanation below). Our infrastructure is currently unable to produce run-time data for the test suite. However, a binary diff shows that the following benchmarks were modified by the change.

Program
---------------------------------------------------------------------------------------
MultiSource/Applications/JM/lencod
MultiSource/Applications/minisat
MultiSource/Benchmarks/Bullet
MultiSource/Benchmarks/MallocBench/espresso
spec2006/h264ref
spec2006/povray

Regarding the lack of performance differences observed in spec2000 and spec2006, the current patch is somewhat limited, and I was planning a follow-on to address the issue. The indices of the GEPs that seed the expressions are forced to be i64. However, the expressions may not require that much precision, so we end up with unneeded extensions and/or narrower vectors than is optimal, and the cost model often prevents us from vectorizing. Please see the test cases in the patch for an example. We are not yet able to vectorize the second case because of this issue.

The follow-on would essentially be to incorporate James's type-shrinking work from the loop vectorizer in order to rewrite the expressions in the narrower type if profitable. Work-in-progress has shown that type-shrinking with this patch can provide a significant performance improvement for spec2006/h264ref, at least.

Please let me know what you think of this plan and whether the optimization is better suited for SLP or SelectionDAG. Thanks again!

[hfinkel]

In D14829#300767, @mssimpso wrote:

Nadav/Hal,

Below are the statistically significant compile-time differences observed for the test suite, spec2000, and spec2006 (there is one). Results were computed from 10 samples each, median aggregation, 95% confidence intervals, and a 0.05 statistical significance level for the Mann–Whitney U test.

Program                                                             Base  Change      %
---------------------------------------------------------------------------------------
MultiSource/Benchmarks/MiBench/security-rijndael/security-rijndael  0.73    0.89  18.88

No statistically significant performance differences were observed for spec2000 and spec2006 on a Cortex-A57-like cpu (but see the explanation below). Our infrastructure is currently unable to produce run-time data for the test suite. However, a binary diff shows that the following benchmarks were modified by the change.

Program
---------------------------------------------------------------------------------------
MultiSource/Applications/JM/lencod
MultiSource/Applications/minisat
MultiSource/Benchmarks/Bullet
MultiSource/Benchmarks/MallocBench/espresso
spec2006/h264ref
spec2006/povray

Regarding the lack of performance differences observed in spec2000 and spec2006, the current patch is somewhat limited, and I was planning a follow-on to address the issue. The indices of the GEPs that seed the expressions are forced to be i64. However, the expressions may not require that much precision, so we end up with unneeded extensions and/or narrower vectors than is optimal, and the cost model often prevents us from vectorizing. Please see the test cases in the patch for an example. We are not yet able to vectorize the second case because of this issue.

The follow-on would essentially be to incorporate James's type-shrinking work from the loop vectorizer in order to rewrite the expressions in the narrower type if profitable. Work-in-progress has shown that type-shrinking with this patch can provide a significant performance improvement for spec2006/h264ref, at least.

Please let me know what you think of this plan and whether the optimization is better suited for SLP or SelectionDAG. Thanks again!

I think SLP is a good place for this, given that the expression trees might be quite large. However, a nearly 20% compile-time slowdown (on an application which surely shows no speedup given your list of binary diffs) is not acceptable. Can you please profile running on that application compilation and propose a patch which limits whatever bad behavior is asserting itself there?

[mssimpso]

Hal,

Thanks very much for the follow-up. Yes, I'm working on compile-time along with the type shrinking work I previously mentioned. I will post an updated revision soon.

[nadav]

In D14829#307505, @mssimpso wrote:

Hal,

Thanks very much for the follow-up. Yes, I'm working on compile-time along with the type shrinking work I previously mentioned. I will post an updated revision soon.

@mssimpso, I don't think that this is the right approach. I think that this should be handled in SelectionDAG, and not in the SLP vectorizer. Your measurements show that there are no performance gains on Spec2000 (and the performance test suite?). Every new heuristic/scan that we add to the SLP vectorizer increases the compile times, and I don't think that adding code to scan loads and start a vec-tree at the load roots is a good heuristic. Many people are using LLVM as a JIT compiler and these people care deeply about compile times. I don't think that in the general case there are many opportunities for vectorization when you search from the address of loads, and I think that your measurements show that this is indeed the case. If you care about the specific pattern of load/store then you can optimize this in SelectionDAG.

[anemet]

@mssimpso, I don't think that this is the right approach. I think that this should be handled in SelectionDAG, and not in the SLP vectorizer. Your measurements show that there are no performance gains on Spec2000 (and the performance test suite?). Every new heuristic/scan that we add to the SLP vectorizer increases the compile times, and I don't think that adding code to scan loads and start a vec-tree at the load roots is a good heuristic. Many people are using LLVM as a JIT compiler and these people care deeply about compile times. I don't think that in the general case there are many opportunities for vectorization when you search from the address of loads, and I think that your measurements show that this is indeed the case. If you care about the specific pattern of load/store then you can optimize this in SelectionDAG.

@nadav, how would the SelectionDAG algorithm be different from what we do here (starting to combine the from load addresses bottom-up)? There is also loop-awareness and being able to vectorize across multiple basic blocks that could be beneficial doing this in LLVM IR.

Also in order to optimize h264ref in SPEC2006int (there is gain in SPEC!) additional type shrinking is needed which is probably easier to do at this level.

lib/Transforms/Vectorize/SLPVectorizer.cpp

415

I think that we typically use the term "vectorization factor"

3149–3152

Hmm, that does not really match the comment for the function. Just because the tree ends in a store we could have a wider load feeding it. Something is not precise enough here.

[nadav]

@anemet In SelectionDAG we can use the memory order chains, and don't have to scan the whole basic block for loads/stores. This is very efficient, and is already used by the load-store merger in SelectionDAG.

I don't think that with the proposed patch we are very likely to catch patterns that cross basic blocks/loops. Can you show examples of real code that we do catch?

[anemet]

In D14829#308530, @nadav wrote:

@anemet In SelectionDAG we can use the memory order chains, and don't have to scan the whole basic block for loads/stores. This is very efficient, and is already used by the load-store merger in SelectionDAG.

But the SLP vectorizer already scans for stores in the entire function so piggybacking on that for loads seems to make sense.

Are you worried about scanning for loads or matching up the "related ones" into the initial bundle? I could see the latter being problematic but I am not sure how the memory order chain could help with that. At the bottom of these chains we have non-consecutive loads (a[b[0] - c[0]], a[b[1] - c[1]], ...). At the top, we have consecutive loads but those should be on independent chains.

[jmolloy]

Resigning from this; Hal and Nadav are reviewing it.

[anemet]

Discussed this more with Nadav in person. Assuming the compile-time issues can be worked out, it probably makes sense to have this in the SLP Vectorizer.

[mssimpso]

Eliminated compile-time regression and addressed Adam's comments.

[mssimpso]

Hi All,

I've updated this patch to eliminate the previously observed compile-time regression in MiBench/security-rijndael. I fixed the issue by moving the contiguous access check (the most expensive check prior to vectorization) to the last step and by processing the accesses in chunks of 16. These two changes follow the existing flow we have for stores to minimze compile-time.

I benchmarked this change together with two others. The additional changes include the type-shrinking work I previously mentioned (D15815) and an additional cost model hook (D15816) to catch the sign extensions we introduce with the type-shrinking. Feel free to comment on those patches as well if interested. Together, the patch set causes no compile-time regressions and improves spec2006/h264ref by ~6% on our Cortex-A57-like architecture.

Thanks!

[Ayal]

It may be better to start from GEPs rather than loads, collecting single-varying-index-GEPs grouped by same-base-pointer that appear in same basic block. That would work for scatters as well as gathers (and a mix ;-), and for other potential users.

If a pair of addresses are consecutive, indeed it would probably be futile to try and vectorize their index computation, because one can be computed from the other rather than independently. But (1) this holds more generally for any simple stride; and (2) it's enough to refrain from treating this pair, why abort all related candidates. Adding "sum += g[2*i]+g[2*i+1];" to your example throws it off needlessly.

+some typos.

lib/Transforms/Vectorize/SLPVectorizer.cpp

415

e[x]pression

3156–3157

We want to base the vector element size on the width of memory operations where possible.

This deserves a comment at the outset, namely that the vector element size is derived from the largest size stored to or loaded from memory (following Adam's comment).

3361

store >> load

3521

typo

[mssimpso]

Addressed Ayal's comments

Thanks very much for the feedback, Ayal! Your suggestions sound reasonable to me. I've updated the patch to seed with GEPs instead of loads and to avoid aborting all candidates if some are consecutive. I also refactored other parts of the patch to match these changes.

I think that starting with GEPs is more straightforward. We were basically already doing this, but only considering the ones used by loads. And I originally chose to abort the candidates if some were consecutive to avoid potential compile-time increases. However, I've benchmarked the new patch with your suggestions and observed no regressions in the LLVM test suite or SPEC. We still bail out early if there are no viable candidates.

[Ayal]

lib/Transforms/Vectorize/SLPVectorizer.cpp

4320

How about simply checking if (isa<SCEVConstant>(Delta)) instead? As said, If one index is easily derived from the other, the two are poor candidates for parallel independent computation. This holds not only when they're exactly adjacent.

Would be good to add tests to make sure this works as intended.

Note that each such index may still be a good candidate, if grouped together with some other, independent indices. Discarding both from consideration may be reasonable, certainly compile-time-wise. The better guidance for such patterns is top-down, as you noted at the beginning.

4321–4323

You could, alternatively, remove correlated GEPs from a set of all current candidates, initialized to all non-nullified GEPs. As long as this set has at-least two remaining candidates.

[mssimpso]

Addressed Ayal's comments. Thanks!

[Ayal]

Functionally this looks fine to me, thanks! I have no further comments. Nadav or Hal should formally LGTMize it though.

Together, the patch set causes no compile-time regressions and improves spec2006/h264ref by ~6% on our Cortex-A57-like architecture.

Cool. It would be good to record other known gains (povray?). This could be useful, e.g., if one attempts to recognize such patterns starting at consecutive loads going downwards, rather than arbitrary independent GEPs going upwards.

[nadav]

LGTM.

[mssimpso]

Thanks very much for the comments everyone! Ayal, in addition to the improvement in spec2006/h264ref (~6%), I observed small improvements in spec2006/povray (~1%) and spec2000/vortex (< 1%) with this patch and the two dependent ones.

[mssimpso]

From Nadav's LGTM.

[lebedev.ri]

llvm/trunk/lib/Transforms/Vectorize/SLPVectorizer.cpp

3538 ↗

(On Diff #44917)

Post-commit review:

I'm 99.9% sure this is wrong.
You are checking the type *of an gep index*,
not the type of the value the gep will produce.
I don't see why you'd care about gep index type.
But you certainly care about resulting type and it's not checked.

You want if (!isValidElementType(GEP->getResultElementType()))

lib/Transforms/Vectorize/SLPVectorizer.cpp

public:
   /// \brief Perform LICM and CSE on the newly generated gather sequences.
   void optimizeGatherSequence();
 
   /// \returns true if it is beneficial to reverse the vector order.
   bool shouldReorder() const {
     return NumLoadsWantToChangeOrder > NumLoadsWantToKeepOrder;
   }
 
+  /// \return The size in bits of the widest memory access in the expression
+  /// tree ending at \p V. This method is used to calculate vector factors.
+  unsigned getMemAccessWidth(Value *V);
+
 private:
   struct TreeEntry;
 
   /// \returns the cost of the vectorizable entry.
   int getEntryCost(TreeEntry *E);
 
   /// This is the recursive part of buildTree.
   void buildTree_rec(ArrayRef<Value *> Roots, unsigned Depth);
while (!ReadyInsts.empty()) {
     NumToSchedule--;
   }
   assert(NumToSchedule == 0 && "could not schedule all instructions");
 
   // Avoid duplicate scheduling of the block.
   BS->ScheduleStart = nullptr;
 }
 
+unsigned BoUpSLP::getMemAccessWidth(Value *V) {
+  auto &DL = F->getParent()->getDataLayout();
+
+  // If V is a store, just return the width of the stored value without
+  // traversing the expression tree. This is the common case.
+  if (auto *Store = dyn_cast<StoreInst>(V))
+    return DL.getTypeSizeInBits(Store->getValueOperand()->getType());
+
+  // If V is not a store, we can traverse the expression tree to find loads
+  // that feed it. The type of the loaded value may indicate a more suitable
+  // width than V's type. We want to base the vector factor on the width of
+  // memory operations where possible.
+  SmallVector<Instruction *, 16> Worklist;
+  SmallPtrSet<Instruction *, 16> Visited;
+  if (auto *I = dyn_cast<Instruction>(V))
+    Worklist.push_back(I);
+
+  // Traverse the expression tree in bottom-up order looking for loads. If we
+  // encounter an instruciton we don't yet handle, we give up.
+  auto MaxWidth = 0u;
+  auto FoundUnknownInst = false;
+  while (!Worklist.empty() && !FoundUnknownInst) {
+    auto *I = Worklist.pop_back_val();
+    Visited.insert(I);
+
+    // We should only be looking at scalar instructions here. If the current
+    // instruction has a vector type, give up.
+    auto *Ty = I->getType();
+    if (isa<VectorType>(Ty))
+      FoundUnknownInst = true;
+
+    // If the current instruction is a load, update MaxWidth to reflect the
+    // width of the loaded value.
+    else if (isa<LoadInst>(I))
+      MaxWidth = std::max(MaxWidth, (unsigned) DL.getTypeSizeInBits(Ty));
+
+    // Otherwise, we need to visit the operands of the instruction. We only
+    // handle the interesting cases from buildTree here. If an operand is an
+    // instruction we haven't yet visited, we add it to the worklist.
+    else if (isa<PHINode>(I) || isa<CastInst>(I) || isa<GetElementPtrInst>(I) ||
+             isa<CmpInst>(I) || isa<SelectInst>(I) || isa<BinaryOperator>(I)) {
+      for (Use &U : I->operands())
+       if (auto *J = dyn_cast<Instruction>(U.get()))
+         if (!Visited.count(J))
+           Worklist.push_back(J);
+    }
+
+    // If we don't yet handle the instruction, give up.
+    else
+      FoundUnknownInst = true;
+  }
+
+  // If we didn't encounter a memory access in the expression tree, or if we
+  // gave up for some reason, just return the width of V.
+  if (!MaxWidth || FoundUnknownInst)
+    return DL.getTypeSizeInBits(V->getType());
+
+  // Otherwise, return the maximum width we found.
+  return MaxWidth;
+}
+
 /// The SLPVectorizer Pass.
 struct SLPVectorizer : public FunctionPass {
   typedef SmallVector<StoreInst *, 8> StoreList;
   typedef MapVector<Value *, StoreList> StoreListMap;
+  typedef SmallVector<WeakVH, 8> LoadList;
+  typedef MapVector<Value *, LoadList> LoadListMap;
 
   /// Pass identification, replacement for typeid
   static char ID;
 
   explicit SLPVectorizer() : FunctionPass(ID) {
     initializeSLPVectorizerPass(*PassRegistry::getPassRegistry());
   }
 
bool runOnFunction(Function &F) override {
     auto *TLIP = getAnalysisIfAvailable<TargetLibraryInfoWrapperPass>();
     TLI = TLIP ? &TLIP->getTLI() : nullptr;
     AA = &getAnalysis<AAResultsWrapperPass>().getAAResults();
     LI = &getAnalysis<LoopInfoWrapperPass>().getLoopInfo();
     DT = &getAnalysis<DominatorTreeWrapperPass>().getDomTree();
     AC = &getAnalysis<AssumptionCacheTracker>().getAssumptionCache(F);
 
     StoreRefs.clear();
+    LoadRefs.clear();
     bool Changed = false;
 
     // If the target claims to have no vector registers don't attempt
     // vectorization.
     if (!TTI->getNumberOfRegisters(true))
       return false;
 
     // Use the vector register size specified by the target unless overridden
bool runOnFunction(Function &F) override {
     // store instructions.
     BoUpSLP R(&F, SE, TTI, TLI, AA, LI, DT, AC);
 
     // A general note: the vectorizer must use BoUpSLP::eraseInstruction() to
     // delete instructions.
 
     // Scan the blocks in the function in post order.
     for (auto BB : post_order(&F.getEntryBlock())) {
+
       // Vectorize trees that end at stores.
-      if (unsigned count = collectStores(BB, R)) {
-        (void)count;
-        DEBUG(dbgs() << "SLP: Found " << count << " stores to vectorize.\n");
+      auto NumStore = collectStores(BB, R);
+      if (NumStore > 0) {
+        DEBUG(dbgs() << "SLP: Found " << NumStore << " stores to vectorize.\n");
         Changed |= vectorizeStoreChains(R);
       }
 
       // Vectorize trees that end at reductions.
       Changed |= vectorizeChainsInBlock(BB, R);
+
+      // Vectorize trees that end at loads. This is intended to catch
+      // gather-like idioms ending at non-consecutive loads.
+      auto NumLoad = collectLoads(BB, R);
+      if (NumLoad > 0) {
+        DEBUG(dbgs() << "SLP: Found " << NumLoad << " loads to vectorize.\n");
+        Changed |= vectorizeLoadChains(BB, R);
+      }
     }
 
     if (Changed) {
       R.optimizeGatherSequence();
       DEBUG(dbgs() << "SLP: vectorized \"" << F.getName() << "\"\n");
       DEBUG(verifyFunction(F));
     }
     return Changed;
void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.addRequired<DominatorTreeWrapperPass>();
     AU.addPreserved<LoopInfoWrapperPass>();
     AU.addPreserved<DominatorTreeWrapperPass>();
     AU.setPreservesCFG();
   }
 
 private:
 
-  /// \brief Collect memory references and sort them according to their base
-  /// object. We sort the stores to their base objects to reduce the cost of the
-  /// quadratic search on the stores. TODO: We can further reduce this cost
-  /// if we flush the chain creation every time we run into a memory barrier.
+  /// \brief Collect stores and sort them according to their base object. We
+  /// sort the stores by their base objects to reduce the cost of the quadratic
+  /// search on the stores. TODO: We can further reduce this cost if we flush
+  /// the chain creation every time we run into a memory barrier.
   unsigned collectStores(BasicBlock *BB, BoUpSLP &R);
 
+  /// \brief Collect loads and sort them according to their base object.
+  unsigned collectLoads(BasicBlock *BB, BoUpSLP &R);
+
   /// \brief Try to vectorize a chain that starts at two arithmetic instrs.
   bool tryToVectorizePair(Value *A, Value *B, BoUpSLP &R);
 
   /// \brief Try to vectorize a list of operands.
   /// \@param BuildVector A list of users to ignore for the purpose of
   ///                     scheduling and that don't need extracting.
   /// \returns true if a value was vectorized.
   bool tryToVectorizeList(ArrayRef<Value *> VL, BoUpSLP &R,
                           ArrayRef<Value *> BuildVector = None,
                           bool allowReorder = false);
 
   /// \brief Try to vectorize a chain that may start at the operands of \V;
   bool tryToVectorize(BinaryOperator *V, BoUpSLP &R);
 
   /// \brief Vectorize the stores that were collected in StoreRefs.
   bool vectorizeStoreChains(BoUpSLP &R);
 
+  /// \brief Vectorize the loads that were collected in LoadRefs.
+  bool vectorizeLoadChains(BasicBlock *BB, BoUpSLP &R);
+
   /// \brief Scan the basic block and look for patterns that are likely to start
   /// a vectorization chain.
   bool vectorizeChainsInBlock(BasicBlock *BB, BoUpSLP &R);
 
   bool vectorizeStoreChain(ArrayRef<Value *> Chain, int CostThreshold,
                            BoUpSLP &R, unsigned VecRegSize);
 
   bool vectorizeStores(ArrayRef<StoreInst *> Stores, int costThreshold,
                        BoUpSLP &R);
 private:
   StoreListMap StoreRefs;
+  LoadListMap LoadRefs;
   unsigned MaxVecRegSize; // This is set by TTI or overridden by cl::opt.
 };
 
 /// \brief Check that the Values in the slice in VL array are still existent in
 /// the WeakVH array.
 /// Vectorization of part of the VL array may cause later values in the VL array
 /// to become invalid. We track when this has happened in the WeakVH array.
 static bool hasValueBeenRAUWed(ArrayRef<Value *> VL, ArrayRef<WeakVH> VH,
                                unsigned SliceBegin, unsigned SliceSize) {
   VL = VL.slice(SliceBegin, SliceSize);
   VH = VH.slice(SliceBegin, SliceSize);
   return !std::equal(VL.begin(), VL.end(), VH.begin());
 }
 
 bool SLPVectorizer::vectorizeStoreChain(ArrayRef<Value *> Chain,
                                         int CostThreshold, BoUpSLP &R,
                                         unsigned VecRegSize) {
   unsigned ChainLen = Chain.size();
   DEBUG(dbgs() << "SLP: Analyzing a store chain of length " << ChainLen
         << "\n");
-  Type *StoreTy = cast<StoreInst>(Chain[0])->getValueOperand()->getType();
-  auto &DL = cast<StoreInst>(Chain[0])->getModule()->getDataLayout();
-  unsigned Sz = DL.getTypeSizeInBits(StoreTy);
+  unsigned Sz = R.getMemAccessWidth(Chain[0]);
   unsigned VF = VecRegSize / Sz;
 
   if (!isPowerOf2_32(Sz) || VF < 2)
     return false;
 
   // Keep track of values that were deleted by vectorizing in the loop below.
   SmallVector<WeakVH, 8> TrackValues(Chain.begin(), Chain.end());
 
for (unsigned Size = MaxVecRegSize; Size >= MinVecRegSize; Size /= 2) {
         break;
       }
     }
   }
 
   return Changed;
 }
 
+unsigned SLPVectorizer::collectLoads(BasicBlock *BB, BoUpSLP &R) {
+  unsigned count = 0;
+  LoadRefs.clear();
+  const DataLayout &DL = BB->getModule()->getDataLayout();
+  for (Instruction &I : *BB) {
+    LoadInst *LI = dyn_cast<LoadInst>(&I);
+    if (!LI)
+      continue;
+
+    // Don't touch volatile loads.
+    if (!LI->isSimple())
+      continue;
+
+    // Check that the pointer points to scalars.
+    Type *Ty = LI->getType();
+    if (!isValidElementType(Ty))
+      continue;
+
+    // Find the base pointer.
+    Value *Ptr = GetUnderlyingObject(LI->getPointerOperand(), DL);
+
+    // Save the load locations.
+    LoadRefs[Ptr].push_back(LI);
+    count++;
+  }
+  return count;
+}
 
 unsigned SLPVectorizer::collectStores(BasicBlock *BB, BoUpSLP &R) {
   unsigned count = 0;
   StoreRefs.clear();
   const DataLayout &DL = BB->getModule()->getDataLayout();
   for (Instruction &I : *BB) {
     StoreInst *SI = dyn_cast<StoreInst>(&I);
     if (!SI)
       continue;
 
     // Don't touch volatile stores.
bool SLPVectorizer::tryToVectorizeList(ArrayRef<Value *> VL, BoUpSLP &R,
   DEBUG(dbgs() << "SLP: Vectorizing a list of length = " << VL.size() << ".\n");
 
   // Check that all of the parts are scalar instructions of the same type.
   Instruction *I0 = dyn_cast<Instruction>(VL[0]);
   if (!I0)
     return false;
 
   unsigned Opcode0 = I0->getOpcode();
-  const DataLayout &DL = I0->getModule()->getDataLayout();
 
-  Type *Ty0 = I0->getType();
-  unsigned Sz = DL.getTypeSizeInBits(Ty0);
   // FIXME: Register size should be a parameter to this function, so we can
   // try different vectorization factors.
+  unsigned Sz = R.getMemAccessWidth(I0);
   unsigned VF = MinVecRegSize / Sz;
 
   for (Value *V : VL) {
     Type *Ty = V->getType();
     if (!isValidElementType(Ty))
       return false;
     Instruction *Inst = dyn_cast<Instruction>(V);
     if (!Inst || Inst->getOpcode() != Opcode0)
if (InsertElementInst *FirstInsertElem = dyn_cast<InsertElementInst>(it)) {
 
       continue;
     }
   }
 
   return Changed;
 }
 
+bool SLPVectorizer::vectorizeLoadChains(BasicBlock *BB, BoUpSLP &R) {
+  auto Changed = false;
+  auto &DL = BB->getModule()->getDataLayout();
+  for (auto &Entry : LoadRefs) {
+
+    // If the load list has fewer than two elements, there's nothing to do.
+    if (Entry.second.size() < 2)
+      continue;
+
+    DEBUG(dbgs() << "SLP: Analyzing a load chain of length "
+                 << Entry.second.size() << ".\n");
+
+    // If the load list accesses contiguous memory, it doesn't make sense to
+    // pursue it further in a bottom-up phase. It should have already been
+    // vectorized if profitable. The code below gives up if a single pair of
+    // consecutive accesses is found.
+    auto IsConsecutive = false;
+    for (unsigned I = 0; I < Entry.second.size() && !IsConsecutive; ++I)
+      for (unsigned J = 0; J < Entry.second.size() && !IsConsecutive; ++J)
+        if (I != J)
+          if (R.isConsecutiveAccess(Entry.second[I], Entry.second[J], DL))
+            IsConsecutive = true;
+    if (IsConsecutive)
+      continue;
+
+    // Try and vectorize trees ending at non-consecutive loads. We are currenly
+    // only interested in gather-like cases of the form:
+    //
+    // ... = g[a[0] - b[0]] + g[a[1] - b[1]] + ...
+    //
+    // where the loads of "a", the loads of "b", and the subtractions can be
+    // performed in parallel. It's likely that detecting this pattern
+    // in a bottom-up phase will be simpler and less costly than building a
+    // full-blown top-down phase beginning at the consecutive loads.
+    SmallVector<Value *, 8> Bundle;
+    for (auto VH : Entry.second) {
+
+      // The load may have been vectorized after we initially collected it. If
+      // so, the WeakVH will have nullified the value.
+      auto *Load = dyn_cast_or_null<LoadInst>(VH);
+      if (!Load)
+        break;
+
+      // We only handle loads in which the pointer operand is a getelementptr.
+      // If we encounter anything else, give up.
+      auto *GEP = dyn_cast<GetElementPtrInst>(Load->getPointerOperand());
+      if (!GEP)
+        break;
+
+      // Furthermore, we only handle getelementptr instructions having a single
+      // non-constant index. Constant indices are not interesting, and we limit
+      // ourselves to getelementptr's with a single index to save compile time.
+      // If we encounter anything else, give up.
+      //
+      // TODO: Relax these restrictions if we find an interesting case that we
+      //       miss and compile time is not an issue.
+      auto *Idx = GEP->idx_begin()->get();
+      if (GEP->getNumIndices() > 1 || isa<Constant>(Idx))
+        break;
+      Bundle.push_back(Idx);
+    }
+
+    // If the bundle has fewer than two elements, there's nothing to do.
+    if (Bundle.size() < 2)
+      continue;
+
+    // Try to vectorize the bundle.
+    Changed |= tryToVectorizeList(Bundle, R);
+  }
+  return Changed;
+}
+
 bool SLPVectorizer::vectorizeStoreChains(BoUpSLP &R) {
   bool Changed = false;
   // Attempt to sort and vectorize each of the store-groups.
   for (StoreListMap::iterator it = StoreRefs.begin(), e = StoreRefs.end();
        it != e; ++it) {
     if (it->second.size() < 2)
       continue;
 

test/Transforms/SLPVectorizer/AArch64/gather-reduce.ll

+; RUN: opt -S -slp-vectorizer -dce -instcombine < %s | FileCheck %s
+
+target datalayout = "e-m:e-i64:64-i128:128-n32:64-S128"
+target triple = "aarch64--linux-gnu"
+
+; These tests check that we vectorize the index calculations in the
+; gather-reduce pattern shown below. We check cases having i32 and i64
+; subtraction.
+;
+; int gather_reduce_8x16(short *a, short *b, short *g, int n) {
+;   int sum = 0;
+;   for (int i = 0; i < n ; ++i) {
+;     sum += g[*a++ - *b++]; sum += g[*a++ - *b++];
+;     sum += g[*a++ - *b++]; sum += g[*a++ - *b++];
+;     sum += g[*a++ - *b++]; sum += g[*a++ - *b++];
+;     sum += g[*a++ - *b++]; sum += g[*a++ - *b++];
+;   }
+;   return sum;
+; }
+
+; CHECK-LABEL: @gather_reduce_8x16_i32
+;
+; CHECK: [[L:%[a-zA-Z0-9.]+]] = load <8 x i16>
+; CHECK: zext <8 x i16> [[L]] to <8 x i32>
+; CHECK: [[S:%[a-zA-Z0-9.]+]] = sub nsw <8 x i32>
+; CHECK: [[X:%[a-zA-Z0-9.]+]] = extractelement <8 x i32> [[S]]
+; CHECK: sext i32 [[X]] to i64
+;
+define i32 @gather_reduce_8x16_i32(i16* nocapture readonly %a, i16* nocapture readonly %b, i16* nocapture readonly %g, i32 %n) {
+entry:
+  %cmp.99 = icmp sgt i32 %n, 0
+  br i1 %cmp.99, label %for.body.preheader, label %for.cond.cleanup
+
+for.body.preheader:
+  br label %for.body
+
+for.cond.cleanup.loopexit:
+  br label %for.cond.cleanup
+
+for.cond.cleanup:
+  %sum.0.lcssa = phi i32 [ 0, %entry ], [ %add66, %for.cond.cleanup.loopexit ]
+  ret i32 %sum.0.lcssa
+
+for.body:
+  %i.0103 = phi i32 [ %inc, %for.body ], [ 0, %for.body.preheader ]
+  %sum.0102 = phi i32 [ %add66, %for.body ], [ 0, %for.body.preheader ]
+  %a.addr.0101 = phi i16* [ %incdec.ptr58, %for.body ], [ %a, %for.body.preheader ]
+  %b.addr.0100 = phi i16* [ %incdec.ptr60, %for.body ], [ %b, %for.body.preheader ]
+  %incdec.ptr = getelementptr inbounds i16, i16* %a.addr.0101, i64 1
+  %0 = load i16, i16* %a.addr.0101, align 2
+  %conv = zext i16 %0 to i32
+  %incdec.ptr1 = getelementptr inbounds i16, i16* %b.addr.0100, i64 1
+  %1 = load i16, i16* %b.addr.0100, align 2
+  %conv2 = zext i16 %1 to i32
+  %sub = sub nsw i32 %conv, %conv2
+  %arrayidx = getelementptr inbounds i16, i16* %g, i32 %sub
+  %2 = load i16, i16* %arrayidx, align 2
+  %conv3 = zext i16 %2 to i32
+  %add = add nsw i32 %conv3, %sum.0102
+  %incdec.ptr4 = getelementptr inbounds i16, i16* %a.addr.0101, i64 2
+  %3 = load i16, i16* %incdec.ptr, align 2
+  %conv5 = zext i16 %3 to i32
+  %incdec.ptr6 = getelementptr inbounds i16, i16* %b.addr.0100, i64 2
+  %4 = load i16, i16* %incdec.ptr1, align 2
+  %conv7 = zext i16 %4 to i32
+  %sub8 = sub nsw i32 %conv5, %conv7
+  %arrayidx10 = getelementptr inbounds i16, i16* %g, i32 %sub8
+  %5 = load i16, i16* %arrayidx10, align 2
+  %conv11 = zext i16 %5 to i32
+  %add12 = add nsw i32 %add, %conv11
+  %incdec.ptr13 = getelementptr inbounds i16, i16* %a.addr.0101, i64 3
+  %6 = load i16, i16* %incdec.ptr4, align 2
+  %conv14 = zext i16 %6 to i32
+  %incdec.ptr15 = getelementptr inbounds i16, i16* %b.addr.0100, i64 3
+  %7 = load i16, i16* %incdec.ptr6, align 2
+  %conv16 = zext i16 %7 to i32
+  %sub17 = sub nsw i32 %conv14, %conv16
+  %arrayidx19 = getelementptr inbounds i16, i16* %g, i32 %sub17
+  %8 = load i16, i16* %arrayidx19, align 2
+  %conv20 = zext i16 %8 to i32
+  %add21 = add nsw i32 %add12, %conv20
+  %incdec.ptr22 = getelementptr inbounds i16, i16* %a.addr.0101, i64 4
+  %9 = load i16, i16* %incdec.ptr13, align 2
+  %conv23 = zext i16 %9 to i32
+  %incdec.ptr24 = getelementptr inbounds i16, i16* %b.addr.0100, i64 4
+  %10 = load i16, i16* %incdec.ptr15, align 2
+  %conv25 = zext i16 %10 to i32
+  %sub26 = sub nsw i32 %conv23, %conv25
+  %arrayidx28 = getelementptr inbounds i16, i16* %g, i32 %sub26
+  %11 = load i16, i16* %arrayidx28, align 2
+  %conv29 = zext i16 %11 to i32
+  %add30 = add nsw i32 %add21, %conv29
+  %incdec.ptr31 = getelementptr inbounds i16, i16* %a.addr.0101, i64 5
+  %12 = load i16, i16* %incdec.ptr22, align 2
+  %conv32 = zext i16 %12 to i32
+  %incdec.ptr33 = getelementptr inbounds i16, i16* %b.addr.0100, i64 5
+  %13 = load i16, i16* %incdec.ptr24, align 2
+  %conv34 = zext i16 %13 to i32
+  %sub35 = sub nsw i32 %conv32, %conv34
+  %arrayidx37 = getelementptr inbounds i16, i16* %g, i32 %sub35
+  %14 = load i16, i16* %arrayidx37, align 2
+  %conv38 = zext i16 %14 to i32
+  %add39 = add nsw i32 %add30, %conv38
+  %incdec.ptr40 = getelementptr inbounds i16, i16* %a.addr.0101, i64 6
+  %15 = load i16, i16* %incdec.ptr31, align 2
+  %conv41 = zext i16 %15 to i32
+  %incdec.ptr42 = getelementptr inbounds i16, i16* %b.addr.0100, i64 6
+  %16 = load i16, i16* %incdec.ptr33, align 2
+  %conv43 = zext i16 %16 to i32
+  %sub44 = sub nsw i32 %conv41, %conv43
+  %arrayidx46 = getelementptr inbounds i16, i16* %g, i32 %sub44
+  %17 = load i16, i16* %arrayidx46, align 2
+  %conv47 = zext i16 %17 to i32
+  %add48 = add nsw i32 %add39, %conv47
+  %incdec.ptr49 = getelementptr inbounds i16, i16* %a.addr.0101, i64 7
+  %18 = load i16, i16* %incdec.ptr40, align 2
+  %conv50 = zext i16 %18 to i32
+  %incdec.ptr51 = getelementptr inbounds i16, i16* %b.addr.0100, i64 7
+  %19 = load i16, i16* %incdec.ptr42, align 2
+  %conv52 = zext i16 %19 to i32
+  %sub53 = sub nsw i32 %conv50, %conv52
+  %arrayidx55 = getelementptr inbounds i16, i16* %g, i32 %sub53
+  %20 = load i16, i16* %arrayidx55, align 2
+  %conv56 = zext i16 %20 to i32
+  %add57 = add nsw i32 %add48, %conv56
+  %incdec.ptr58 = getelementptr inbounds i16, i16* %a.addr.0101, i64 8
+  %21 = load i16, i16* %incdec.ptr49, align 2
+  %conv59 = zext i16 %21 to i32
+  %incdec.ptr60 = getelementptr inbounds i16, i16* %b.addr.0100, i64 8
+  %22 = load i16, i16* %incdec.ptr51, align 2
+  %conv61 = zext i16 %22 to i32
+  %sub62 = sub nsw i32 %conv59, %conv61
+  %arrayidx64 = getelementptr inbounds i16, i16* %g, i32 %sub62
+  %23 = load i16, i16* %arrayidx64, align 2
+  %conv65 = zext i16 %23 to i32
+  %add66 = add nsw i32 %add57, %conv65
+  %inc = add nuw nsw i32 %i.0103, 1
+  %exitcond = icmp eq i32 %inc, %n
+  br i1 %exitcond, label %for.cond.cleanup.loopexit, label %for.body
+}
+
+; CHECK-LABEL: @gather_reduce_8x16_i64
+;
+; CHECK-NOT: load <8 x i16>
+;
+; FIXME: We are currently unable to vectorize the case with i64 subtraction
+;        because the zero extensions are too expensive. The solution here is to
+;        convert the i64 subtractions to i32 subtractions during vectorization.
+;        This would then match the case above.
+;
+define i32 @gather_reduce_8x16_i64(i16* nocapture readonly %a, i16* nocapture readonly %b, i16* nocapture readonly %g, i32 %n) {
+entry:
+  %cmp.99 = icmp sgt i32 %n, 0
+  br i1 %cmp.99, label %for.body.preheader, label %for.cond.cleanup
+
+for.body.preheader:
+  br label %for.body
+
+for.cond.cleanup.loopexit:
+  br label %for.cond.cleanup
+
+for.cond.cleanup:
+  %sum.0.lcssa = phi i32 [ 0, %entry ], [ %add66, %for.cond.cleanup.loopexit ]
+  ret i32 %sum.0.lcssa
+
+for.body:
+  %i.0103 = phi i32 [ %inc, %for.body ], [ 0, %for.body.preheader ]
+  %sum.0102 = phi i32 [ %add66, %for.body ], [ 0, %for.body.preheader ]
+  %a.addr.0101 = phi i16* [ %incdec.ptr58, %for.body ], [ %a, %for.body.preheader ]
+  %b.addr.0100 = phi i16* [ %incdec.ptr60, %for.body ], [ %b, %for.body.preheader ]
+  %incdec.ptr = getelementptr inbounds i16, i16* %a.addr.0101, i64 1
+  %0 = load i16, i16* %a.addr.0101, align 2
+  %conv = zext i16 %0 to i64
+  %incdec.ptr1 = getelementptr inbounds i16, i16* %b.addr.0100, i64 1
+  %1 = load i16, i16* %b.addr.0100, align 2
+  %conv2 = zext i16 %1 to i64
+  %sub = sub nsw i64 %conv, %conv2
+  %arrayidx = getelementptr inbounds i16, i16* %g, i64 %sub
+  %2 = load i16, i16* %arrayidx, align 2
+  %conv3 = zext i16 %2 to i32
+  %add = add nsw i32 %conv3, %sum.0102
+  %incdec.ptr4 = getelementptr inbounds i16, i16* %a.addr.0101, i64 2
+  %3 = load i16, i16* %incdec.ptr, align 2
+  %conv5 = zext i16 %3 to i64
+  %incdec.ptr6 = getelementptr inbounds i16, i16* %b.addr.0100, i64 2
+  %4 = load i16, i16* %incdec.ptr1, align 2
+  %conv7 = zext i16 %4 to i64
+  %sub8 = sub nsw i64 %conv5, %conv7
+  %arrayidx10 = getelementptr inbounds i16, i16* %g, i64 %sub8
+  %5 = load i16, i16* %arrayidx10, align 2
+  %conv11 = zext i16 %5 to i32
+  %add12 = add nsw i32 %add, %conv11
+  %incdec.ptr13 = getelementptr inbounds i16, i16* %a.addr.0101, i64 3
+  %6 = load i16, i16* %incdec.ptr4, align 2
+  %conv14 = zext i16 %6 to i64
+  %incdec.ptr15 = getelementptr inbounds i16, i16* %b.addr.0100, i64 3
+  %7 = load i16, i16* %incdec.ptr6, align 2
+  %conv16 = zext i16 %7 to i64
+  %sub17 = sub nsw i64 %conv14, %conv16
+  %arrayidx19 = getelementptr inbounds i16, i16* %g, i64 %sub17
+  %8 = load i16, i16* %arrayidx19, align 2
+  %conv20 = zext i16 %8 to i32
+  %add21 = add nsw i32 %add12, %conv20
+  %incdec.ptr22 = getelementptr inbounds i16, i16* %a.addr.0101, i64 4
+  %9 = load i16, i16* %incdec.ptr13, align 2
+  %conv23 = zext i16 %9 to i64
+  %incdec.ptr24 = getelementptr inbounds i16, i16* %b.addr.0100, i64 4
+  %10 = load i16, i16* %incdec.ptr15, align 2
+  %conv25 = zext i16 %10 to i64
+  %sub26 = sub nsw i64 %conv23, %conv25
+  %arrayidx28 = getelementptr inbounds i16, i16* %g, i64 %sub26
+  %11 = load i16, i16* %arrayidx28, align 2
+  %conv29 = zext i16 %11 to i32
+  %add30 = add nsw i32 %add21, %conv29
+  %incdec.ptr31 = getelementptr inbounds i16, i16* %a.addr.0101, i64 5
+  %12 = load i16, i16* %incdec.ptr22, align 2
+  %conv32 = zext i16 %12 to i64
+  %incdec.ptr33 = getelementptr inbounds i16, i16* %b.addr.0100, i64 5
+  %13 = load i16, i16* %incdec.ptr24, align 2
+  %conv34 = zext i16 %13 to i64
+  %sub35 = sub nsw i64 %conv32, %conv34
+  %arrayidx37 = getelementptr inbounds i16, i16* %g, i64 %sub35
+  %14 = load i16, i16* %arrayidx37, align 2
+  %conv38 = zext i16 %14 to i32
+  %add39 = add nsw i32 %add30, %conv38
+  %incdec.ptr40 = getelementptr inbounds i16, i16* %a.addr.0101, i64 6
+  %15 = load i16, i16* %incdec.ptr31, align 2
+  %conv41 = zext i16 %15 to i64
+  %incdec.ptr42 = getelementptr inbounds i16, i16* %b.addr.0100, i64 6
+  %16 = load i16, i16* %incdec.ptr33, align 2
+  %conv43 = zext i16 %16 to i64
+  %sub44 = sub nsw i64 %conv41, %conv43
+  %arrayidx46 = getelementptr inbounds i16, i16* %g, i64 %sub44
+  %17 = load i16, i16* %arrayidx46, align 2
+  %conv47 = zext i16 %17 to i32
+  %add48 = add nsw i32 %add39, %conv47
+  %incdec.ptr49 = getelementptr inbounds i16, i16* %a.addr.0101, i64 7
+  %18 = load i16, i16* %incdec.ptr40, align 2
+  %conv50 = zext i16 %18 to i64
+  %incdec.ptr51 = getelementptr inbounds i16, i16* %b.addr.0100, i64 7
+  %19 = load i16, i16* %incdec.ptr42, align 2
+  %conv52 = zext i16 %19 to i64
+  %sub53 = sub nsw i64 %conv50, %conv52
+  %arrayidx55 = getelementptr inbounds i16, i16* %g, i64 %sub53
+  %20 = load i16, i16* %arrayidx55, align 2
+  %conv56 = zext i16 %20 to i32
+  %add57 = add nsw i32 %add48, %conv56
+  %incdec.ptr58 = getelementptr inbounds i16, i16* %a.addr.0101, i64 8
+  %21 = load i16, i16* %incdec.ptr49, align 2
+  %conv59 = zext i16 %21 to i64
+  %incdec.ptr60 = getelementptr inbounds i16, i16* %b.addr.0100, i64 8
+  %22 = load i16, i16* %incdec.ptr51, align 2
+  %conv61 = zext i16 %22 to i64
+  %sub62 = sub nsw i64 %conv59, %conv61
+  %arrayidx64 = getelementptr inbounds i16, i16* %g, i64 %sub62
+  %23 = load i16, i16* %arrayidx64, align 2
+  %conv65 = zext i16 %23 to i32
+  %add66 = add nsw i32 %add57, %conv65
+  %inc = add nuw nsw i32 %i.0103, 1
+  %exitcond = icmp eq i32 %inc, %n
+  br i1 %exitcond, label %for.cond.cleanup.loopexit, label %for.body
+}