Extend LoopVectorizationLegality::isConsecutivePtr to handle multiple level GEPs

[wmi]

Summary

Find another testcase requiring isConsecutivePtr to handle multiple level GEPs (described in https://llvm.org/bugs/show_bug.cgi?id=23580). Drop D10281 and try a simpler patch to solve it.

[anemet]

Did you see D20789? (The commit is currently reverted due to some failures.)

[wmi]

Ah, I didn't see the patch. Thanks for letting me know. I just verified the patch solved the problem for the testcases I cared about, so I will dropped this patch and waited D20789 to be recommitted.

Thanks,
Wei.

include/llvm/Analysis/VectorUtils.h

 /// \brief Returns intrinsic ID for call.
 /// For the input call instruction it finds mapping intrinsic and returns
 /// its intrinsic ID, in case it does not found it return not_intrinsic.
 Intrinsic::ID getVectorIntrinsicIDForCall(const CallInst *CI,
                                           const TargetLibraryInfo *TLI);
 
 /// \brief Find the operand of the GEP that should be checked for consecutive
 /// stores. This ignores trailing indices that have no effect on the final
-/// pointer.
-unsigned getGEPInductionOperand(const GetElementPtrInst *Gep);
+/// pointer. If \p Gep is not the last level of GEP computing store address,
+/// the mem access type of store -- \p MemAccessTy should be used instead
+/// of Gep's result type.
+unsigned getGEPInductionOperand(const GetElementPtrInst *Gep,
+                                Type *MemAccessTy = nullptr);
 
 /// \brief If the argument is a GEP, then returns the operand identified by
 /// getGEPInductionOperand. However, if there is some other non-loop-invariant
 /// operand, it returns that instead.
 Value *stripGetElementPtr(Value *Ptr, ScalarEvolution *SE, Loop *Lp);
 
 /// \brief If a value has only one user that is a CastInst, return it.
 Value *getUniqueCastUse(Value *Ptr, Loop *Lp, Type *Ty);

lib/Analysis/VectorUtils.cpp

if (isTriviallyVectorizable(ID) || ID == Intrinsic::lifetime_start ||
       ID == Intrinsic::lifetime_end || ID == Intrinsic::assume)
     return ID;
   return Intrinsic::not_intrinsic;
 }
 
 /// \brief Find the operand of the GEP that should be checked for consecutive
 /// stores. This ignores trailing indices that have no effect on the final
 /// pointer.
-unsigned llvm::getGEPInductionOperand(const GetElementPtrInst *Gep) {
+unsigned llvm::getGEPInductionOperand(const GetElementPtrInst *Gep,
+                                      Type *MemAccessTy) {
   const DataLayout &DL = Gep->getModule()->getDataLayout();
   unsigned LastOperand = Gep->getNumOperands() - 1;
-  unsigned GEPAllocSize = DL.getTypeAllocSize(Gep->getResultElementType());
+  unsigned GEPAllocSize = DL.getTypeAllocSize(
+      MemAccessTy ? MemAccessTy : Gep->getResultElementType());
 
   // Walk backwards and try to peel off zeros.
   while (LastOperand > 1 && match(Gep->getOperand(LastOperand), m_Zero())) {
     // Find the type we're currently indexing into.
     gep_type_iterator GEPTI = gep_type_begin(Gep);
     std::advance(GEPTI, LastOperand - 1);
 
     // If it's a type with the same allocation size as the result of the GEP we

lib/Transforms/Vectorize/LoopVectorize.cpp

public:
   /// Check if this pointer is consecutive when vectorizing. This happens
   /// when the last index of the GEP is the induction variable, or that the
   /// pointer itself is an induction variable.
   /// This check allows us to vectorize A[idx] into a wide load/store.
   /// Returns:
   /// 0 - Stride is unknown or non-consecutive.
   /// 1 - Address is consecutive.
   /// -1 - Address is consecutive, and decreasing.
-  int isConsecutivePtr(Value *Ptr);
+  int isConsecutivePtr(Value *Ptr, Type *MemAccessTy = nullptr);
 
   /// Returns true if the value V is uniform within the loop.
   bool isUniform(Value *V);
 
   /// Returns true if this instruction will remain scalar after vectorization.
   bool isUniformAfterVectorization(Instruction *I) { return Uniforms.count(I); }
 
   /// Returns the information that we collected about runtime memory check.
Value *InnerLoopVectorizer::getStepVector(Value *Val, int StartIdx,
   Step = Builder.CreateVectorSplat(VLen, Step);
   assert(Step->getType() == Val->getType() && "Invalid step vec");
   // FIXME: The newly created binary instructions should contain nsw/nuw flags,
   // which can be found from the original scalar operations.
   Step = Builder.CreateMul(Cv, Step);
   return Builder.CreateAdd(Val, Step, "induction");
 }
 
-int LoopVectorizationLegality::isConsecutivePtr(Value *Ptr) {
+int LoopVectorizationLegality::isConsecutivePtr(Value *Ptr, Type *MemAccessTy) {
   assert(Ptr->getType()->isPointerTy() && "Unexpected non-ptr");
   auto *SE = PSE.getSE();
-  // Make sure that the pointer does not point to structs.
-  if (Ptr->getType()->getPointerElementType()->isAggregateType())
+  // If MemAccessTy is nullptr, assuming we are checking a Ptr directly used
+  // as load/store address, so the pointer cannot point to structs.
+  if (!MemAccessTy &&
+      Ptr->getType()->getPointerElementType()->isAggregateType())
+    return 0;
+
+  // If MemAccessTy is not nullptr, we are checking a Ptr indirectly used as
+  // load/store address.
+  // Make sure Ptr's element type has the same size as MemAccessTy, so if only
+  // Ptr is consecutive, it is consecutive relative to MemAccessTy.
+  auto &DL = TheLoop->getHeader()->getModule()->getDataLayout();
+  if (MemAccessTy &&
+      DL.getTypeAllocSize(MemAccessTy) !=
+          DL.getTypeAllocSize(Ptr->getType()->getPointerElementType()))
     return 0;
 
   // If this value is a pointer induction variable, we know it is consecutive.
   PHINode *Phi = dyn_cast_or_null<PHINode>(Ptr);
   if (Phi && Inductions.count(Phi)) {
     InductionDescriptor II = Inductions[Phi];
     return II.getConsecutiveDirection();
   }
 
   GetElementPtrInst *Gep = getGEPInstruction(Ptr);
   if (!Gep)
     return 0;
 
   unsigned NumOperands = Gep->getNumOperands();
   Value *GpPtr = Gep->getPointerOperand();
   // If this GEP value is a consecutive pointer induction variable and all of
   // the indices are constant, then we know it is consecutive.
   Phi = dyn_cast<PHINode>(GpPtr);
   if (Phi && Inductions.count(Phi)) {
 
     // Make sure that the pointer does not point to structs.
     PointerType *GepPtrType = cast<PointerType>(GpPtr->getType());
-    if (GepPtrType->getElementType()->isAggregateType())
+    if (GepPtrType->getElementType()->isAggregateType()) {
       return 0;
+    }
 
     // Make sure that all of the index operands are loop invariant.
     for (unsigned i = 1; i < NumOperands; ++i)
       if (!SE->isLoopInvariant(PSE.getSCEV(Gep->getOperand(i)), TheLoop))
         return 0;
 
     InductionDescriptor II = Inductions[Phi];
     return II.getConsecutiveDirection();
   }
 
-  unsigned InductionOperand = getGEPInductionOperand(Gep);
+  // If GpPtr is a consecutive pointer relative to MemAccessTy and
+  // all of the GEP indices are constant, then we know it is consecutive.
+  if (int Direction = isConsecutivePtr(
+          GpPtr, MemAccessTy ? MemAccessTy : Gep->getResultElementType())) {
+    // Make sure that all of the index operands are loop invariant.
+    for (unsigned i = 1; i < NumOperands; ++i)
+      if (!SE->isLoopInvariant(PSE.getSCEV(Gep->getOperand(i)), TheLoop))
+        return 0;
+    return Direction;
+  }
+
+  unsigned InductionOperand = getGEPInductionOperand(Gep, MemAccessTy);
 
   // Check that all of the gep indices are uniform except for our induction
   // operand.
   for (unsigned i = 0; i != NumOperands; ++i)
     if (i != InductionOperand &&
         !SE->isLoopInvariant(PSE.getSCEV(Gep->getOperand(i)), TheLoop))
       return 0;
 
void InnerLoopVectorizer::vectorizeMemoryInstruction(Instruction *Instr) {
 
   Constant *Zero = Builder.getInt32(0);
   VectorParts &Entry = WidenMap.get(Instr);
   VectorParts VectorGep;
 
   // Handle consecutive loads/stores.
   GetElementPtrInst *Gep = getGEPInstruction(Ptr);
   if (ConsecutiveStride) {
-    if (Gep && Legal->isInductionVariable(Gep->getPointerOperand())) {
+    if (Gep &&
+        Legal->isConsecutivePtr(Gep->getPointerOperand(),
+                                Gep->getResultElementType())) {
       setDebugLocFromInst(Builder, Gep);
       Value *PtrOperand = Gep->getPointerOperand();
       Value *FirstBasePtr = getVectorValue(PtrOperand)[0];
       FirstBasePtr = Builder.CreateExtractElement(FirstBasePtr, Zero);
 
       // Create the new GEP with the new induction variable.
       GetElementPtrInst *Gep2 = cast<GetElementPtrInst>(Gep->clone());
       Gep2->setOperand(0, FirstBasePtr);

test/Transforms/LoopVectorize/consecutive-ptr.ll

+; RUN: opt < %s -loop-vectorize -force-vector-interleave=1 -force-vector-width=4 -S | FileCheck %s
+
+; Make sure loop vectorizer knows the store address is consecutive and generates wide store.
+; CHECK-LABEL: @foo(
+; CHECK: store <4 x i32>
+
+define void @foo(i32 %N, i32* nocapture %ps) #0 {
+entry:
+  %cmp6 = icmp sgt i32 %N, 0
+  br i1 %cmp6, label %for.body.preheader, label %for.end
+
+for.body.preheader:                               ; preds = %entry
+  br label %for.body
+
+for.body:                                         ; preds = %for.body.preheader, %for.body
+  %indvars.iv = phi i64 [ %indvars.iv.next, %for.body ], [ 0, %for.body.preheader ]
+  %add.ptr = getelementptr inbounds i32, i32* %ps, i64 %indvars.iv
+  %add.ptr1 = getelementptr inbounds i32, i32* %add.ptr, i64 1
+  %tmp0 = trunc i64 %indvars.iv to i32
+  store i32 %tmp0, i32* %add.ptr1, align 4
+  %indvars.iv.next = add nuw nsw i64 %indvars.iv, 1
+  %lftr.wideiv = trunc i64 %indvars.iv.next to i32
+  %exitcond = icmp eq i32 %lftr.wideiv, %N
+  br i1 %exitcond, label %for.end.loopexit, label %for.body
+
+for.end.loopexit:                                 ; preds = %for.body
+  br label %for.end
+
+for.end:                                          ; preds = %for.end.loopexit, %entry
+  ret void
+}
+
+; Make sure loop vectorizer knows the store address is consecutive and generates wide store.
+; CHECK-LABEL: @goo(
+; CHECK: store <4 x i32>
+
+%struct.B = type { i32 }
+define void @goo(i32 %N, %struct.B* %ps) #0 {
+entry:
+  %cmp6 = icmp sgt i32 %N, 0
+  br i1 %cmp6, label %for.body.preheader, label %for.end
+
+for.body.preheader:                               ; preds = %entry
+  br label %for.body
+
+for.body:                                         ; preds = %for.body.preheader, %for.body
+  %indvars.iv = phi i64 [ %indvars.iv.next, %for.body ], [ 0, %for.body.preheader ]
+  %arrayidx16 = getelementptr inbounds %struct.B, %struct.B* %ps, i64 %indvars.iv
+  %ival = getelementptr inbounds %struct.B, %struct.B* %arrayidx16, i64 0, i32 0
+  %tmp1 = trunc i64 %indvars.iv to i32
+  store i32 %tmp1, i32* %ival, align 4
+  %indvars.iv.next = add nuw nsw i64 %indvars.iv, 1
+  %lftr.wideiv = trunc i64 %indvars.iv.next to i32
+  %exitcond = icmp eq i32 %lftr.wideiv, %N
+  br i1 %exitcond, label %for.end.loopexit, label %for.body
+
+for.end.loopexit:                                 ; preds = %for.body
+  br label %for.end
+
+for.end:                                          ; preds = %for.end.loopexit, %entry
+  ret void
+}
+