1 files changed, 31 insertions, 12 deletions

diff --git a/lib/Transforms/Vectorize/SLPVectorizer.cpp b/lib/Transforms/Vectorize/SLPVectorizer.cpp
index f0a5e45..6498784 100644
--- a/lib/Transforms/Vectorize/SLPVectorizer.cpp
+++ b/lib/Transforms/Vectorize/SLPVectorizer.cpp
@@ -984,31 +984,50 @@ bool BoUpSLP::isConsecutiveAccess(Value *A, Value *B) {
     return false;
 
   // Calculate a constant offset from the base pointer without using SCEV
-  // in the supported cases. 
+  // in the supported cases.
   // TODO: Add support for the case where one of the pointers is a GEP that
   // uses the other pointer.
   GetElementPtrInst *GepA = dyn_cast<GetElementPtrInst>(PtrA);
   GetElementPtrInst *GepB = dyn_cast<GetElementPtrInst>(PtrB);
-  if (GepA && GepB && GepA->getPointerOperand() == GepB->getPointerOperand()) {
-    unsigned BW = DL->getPointerSizeInBits(ASA);
-    APInt OffsetA(BW, 0) ,OffsetB(BW, 0);
 
+  unsigned BW = DL->getPointerSizeInBits(ASA);
+  Type *Ty = cast<PointerType>(PtrA->getType())->getElementType();
+  int64_t Sz = DL->getTypeStoreSize(Ty);
+
+  // If both pointers are GEPs:
+  if (GepA && GepB) {
+    // Check that they have the same base pointer.
+    if (GepA->getPointerOperand() != GepB->getPointerOperand())
+      return false;
+
+    // Check if the geps use a constant offset.
+    APInt OffsetA(BW, 0) ,OffsetB(BW, 0);
     if (GepA->accumulateConstantOffset(*DL, OffsetA) &&
-        GepB->accumulateConstantOffset(*DL, OffsetB)) {
-      Type *Ty = cast<PointerType>(PtrA->getType())->getElementType();
-      int64_t Sz = DL->getTypeStoreSize(Ty);
+        GepB->accumulateConstantOffset(*DL, OffsetB))
       return ((OffsetB.getSExtValue() - OffsetA.getSExtValue()) == Sz);
+
+    // Try to strip the geps. This makes SCEV faster.
+    if (GepA->getNumIndices() == 1 && GepB->getNumIndices() == 1) {
+      PtrA = GepA->getOperand(1);
+      PtrB = GepB->getOperand(1);
+      Sz = 1;
     }
   }
 
+  // Check if PtrA is the base and PtrB is a constant offset.
+  if (GepB && GepB->getPointerOperand() == PtrA) {
+    APInt Offset(BW, 0);
+    if (GepB->accumulateConstantOffset(*DL, Offset))
+      return Offset.getZExtValue() == DL->getTypeStoreSize(Ty);
+  }
+
+  // GepA can't use PtrB as a base pointer.
+  if (GepA && GepA->getPointerOperand() == PtrB)
+    return false;
+
   // Calculate the distance.
   const SCEV *PtrSCEVA = SE->getSCEV(PtrA);
   const SCEV *PtrSCEVB = SE->getSCEV(PtrB);
-  Type *Ty = cast<PointerType>(PtrA->getType())->getElementType();
-  // The instructions are consecutive if the size of the first load/store is
-  // the same as the offset.
-  int64_t Sz = DL->getTypeStoreSize(Ty);
-
   const SCEV *C = SE->getConstant(PtrSCEVA->getType(), Sz);
   const SCEV *X = SE->getAddExpr(PtrSCEVA, C);
   return X == PtrSCEVB;