[LoopVectorize] Classify pointer induction updates as scalar only if they have one use

[kmclaughlin]

Summary

collectLoopScalars collects pointer induction updates in ScalarPtrs, assuming
that the instruction will be scalar after vectorization. This may crash later
in VPReplicateRecipe::execute() if there there is another user of the instruction
other than the Phi node which needs to be widened.

This changes collectLoopScalars so that if there are any other users of
Update other than a Phi node, it is not added to ScalarPtrs.

[david-arm]

This looks like a sensible fix to me. I just had a couple of minor comments ...

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

5134

Hi @kmclaughlin, I think it's probably better to use isa<PHINode> here instead of dyn_cast<PHINode>.

llvm/test/Transforms/LoopVectorize/AArch64/sve-widen-gep.ll

121

Maybe it's worth simplifying the test a little by killing off the CHECK lines before and after the vector.body just to make it easier to read?

[kmclaughlin]

• Replaced use of dyn_cast<PHINode> in collectLoopScalars with isa<PHINode>
• Reduced the CHECK lines before & after vector.body for the new test added to sve-widen-gep.ll

[david-arm]

LGTM. Thanks for making the changes!

[fhahn]

I think it would be good to update the title to better reflect the change here. From the title it sounds like a change in VPReplicateRecipe::execute() , but it actually adjusts how pointers are classified which as a side-effect fixes a crash.

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

5133

nit: 'llvm:: should not be needed. Can you also add comment here to explain the reasoning?

Also, no need to capture everything in the lambda.

llvm/test/Transforms/LoopVectorize/AArch64/sve-widen-gep.ll

166

just return void?

[fhahn]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

5137

Hm, with this change we might miss some cases where the uses of Update are still scalar. Would it be better to only return early here if the only user of the update instruction is the pointer phi and let the code below handle the other cases?

[kmclaughlin]

• Return early from collectLoopScalars if Update has one use. Removed line to insert Update into PossibleNonScalarPtrs if this is not the case to instead let the code below the if (isScalarPtrInduction... block handle other cases.
• Return void from @pointer_induction test

[david-arm]

LGTM! It looks like you've addressed @fhahn's comments.

[fhahn]

LGTM, thanks for the update. I added a few more tiny suggestions to slightly improve readability of the test case that would be good to include.

llvm/test/Transforms/LoopVectorize/AArch64/sve-widen-gep.ll

151

nit: could just use a i8* %start argument instead of [64 x i8]?

152

nit: unused pointer?

156

nit: might be good to drop the .sroa prefix, perhaps just something like %ptr.phi or something and then %ptr.phi.next instead of %incdec.ptr.

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

void LoopVectorizationCostModel::collectLoopScalars(ElementCount VF) {
   auto evaluatePtrUse = [&](Instruction *MemAccess, Value *Ptr) {
     if (isScalarPtrInduction(MemAccess, Ptr)) {
       Worklist.insert(cast<Instruction>(Ptr));
       LLVM_DEBUG(dbgs() << "LV: Found new scalar instruction: " << *Ptr
                         << "\n");
 
       Instruction *Update = cast<Instruction>(
           cast<PHINode>(Ptr)->getIncomingValueForBlock(Latch));
+
+      // If there is more than one user of Update (Ptr), we shouldn't assume it
+      // will be scalar after vectorisation as other users of the instruction
+      // may require widening. Otherwise, add it to ScalarPtrs.
+      if (Update->hasOneUse() && cast<Value>(*Update->user_begin()) == Ptr) {
         ScalarPtrs.insert(Update);
         return;
       }
+    }
     // We only care about bitcast and getelementptr instructions contained in
     // the loop.
     if (!isLoopVaryingBitCastOrGEP(Ptr))
       return;
 
     // If the pointer has already been identified as scalar (e.g., if it was
     // also identified as uniform), there's nothing to do.
     auto *I = cast<Instruction>(Ptr);

llvm/test/Transforms/LoopVectorize/AArch64/sve-widen-gep.ll

loop.body:                                    ; preds = %loop.body, %entry
   %iv.next = add nuw i64 %iv, 1
   %c = icmp ne i64 %iv.next, %N
   br i1 %c, label %loop.body, label %exit, !llvm.loop !0
 
 exit:                            ; preds = %loop.body
   ret void
 }
 
+define void @pointer_induction(i8* noalias %start, i64 %N) {
+; CHECK-LABEL: @pointer_induction(
+; CHECK:       vector.ph:
+; CHECK:         [[BROADCAST_SPLATINSERT:%.*]] = insertelement <vscale x 2 x i8*> poison, i8* [[START:%.*]], i32 0
+; CHECK-NEXT:    [[BROADCAST_SPLAT:%.*]] = shufflevector <vscale x 2 x i8*> [[BROADCAST_SPLATINSERT]], <vscale x 2 x i8*> poison, <vscale x 2 x i32> zeroinitializer
+; CHECK:       vector.body:
+; CHECK-NEXT:    [[INDEX1:%.*]] = phi i64 [ 0, [[VECTOR_PH]] ], [ [[INDEX_NEXT:%.*]], [[VECTOR_BODY]] ]
+; CHECK-NEXT:    [[TMP5:%.*]] = call <vscale x 2 x i64> @llvm.experimental.stepvector.nxv2i64()
+; CHECK-NEXT:    [[DOTSPLATINSERT:%.*]] = insertelement <vscale x 2 x i64> poison, i64 [[INDEX1]], i32 0
+; CHECK-NEXT:    [[DOTSPLAT:%.*]] = shufflevector <vscale x 2 x i64> [[DOTSPLATINSERT]], <vscale x 2 x i64> poison, <vscale x 2 x i32> zeroinitializer
+; CHECK-NEXT:    [[TMP6:%.*]] = add <vscale x 2 x i64> shufflevector (<vscale x 2 x i64> insertelement (<vscale x 2 x i64> poison, i64 0, i32 0), <vscale x 2 x i64> poison, <vscale x 2 x i32> zeroinitializer), [[TMP5]]
+; CHECK-NEXT:    [[TMP7:%.*]] = add <vscale x 2 x i64> [[DOTSPLAT]], [[TMP6]]
+; CHECK-NEXT:    [[NEXT_GEP:%.*]] = getelementptr i8, i8* [[START]], <vscale x 2 x i64> [[TMP7]]
+; CHECK-NEXT:    [[TMP8:%.*]] = add i64 [[INDEX1]], 0
+; CHECK-NEXT:    [[NEXT_GEP3:%.*]] = getelementptr i8, i8* [[START]], i64 [[TMP8]]
+; CHECK-NEXT:    [[TMP9:%.*]] = add i64 [[INDEX1]], 1
+; CHECK-NEXT:    [[NEXT_GEP4:%.*]] = getelementptr i8, i8* [[START]], i64 [[TMP9]]
+; CHECK-NEXT:    [[TMP10:%.*]] = add i64 [[INDEX1]], 0
+; CHECK-NEXT:    [[TMP11:%.*]] = getelementptr i8, i8* [[NEXT_GEP3]], i32 0
+; CHECK-NEXT:    [[TMP12:%.*]] = bitcast i8* [[TMP11]] to <vscale x 2 x i8>*
+; CHECK-NEXT:    [[WIDE_LOAD:%.*]] = load <vscale x 2 x i8>, <vscale x 2 x i8>* [[TMP12]], align 1
+; CHECK-NEXT:    [[TMP13:%.*]] = getelementptr inbounds i8, <vscale x 2 x i8*> [[NEXT_GEP]], i64 1
+; CHECK-NEXT:    [[TMP14:%.*]] = icmp eq <vscale x 2 x i8*> [[TMP13]], [[BROADCAST_SPLAT]]
+; CHECK-NEXT:    [[TMP15:%.*]] = call i64 @llvm.vscale.i64()
+; CHECK-NEXT:    [[TMP16:%.*]] = mul i64 [[TMP15]], 2
+; CHECK-NEXT:    [[INDEX_NEXT]] = add nuw i64 [[INDEX1]], [[TMP16]]
+; CHECK-NEXT:    [[TMP17:%.*]] = icmp eq i64 [[INDEX_NEXT]], [[N_VEC]]
+; CHECK-NEXT:    br i1 [[TMP17]], label [[MIDDLE_BLOCK:%.*]], label [[VECTOR_BODY]], !llvm.loop [[LOOP4:![0-9]+]]
+;
+entry:
+  br label %for.body
+
+for.body:
+  %ptr.phi = phi i8* [ %ptr.phi.next, %for.body ], [ %start, %entry ]
+  %index = phi i64 [ %index_nxt, %for.body ], [ 0, %entry ]
+  %index_nxt = add i64 %index, 1
+  %0 = load i8, i8* %ptr.phi, align 1
+  %ptr.phi.next = getelementptr inbounds i8, i8* %ptr.phi, i64 1
+  %cmp.i.not = icmp eq i8* %ptr.phi.next, %start
+  %cmp = icmp ult i64 %index, %N
+  br i1 %cmp, label %for.body, label %end, !llvm.loop !0
+
+end:
+  ret void
+}
 
 attributes #0 = {"target-features"="+sve"}
 
 !0 = distinct !{!0, !1, !2, !3}
 !1 = !{!"llvm.loop.interleave.count", i32 1}
 !2 = !{!"llvm.loop.vectorize.width", i32 2}
 !3 = !{!"llvm.loop.vectorize.scalable.enable", i1 true}
 !4 = !{ !5 }
 !5 = distinct !{ !5, !6 }
 !6 = distinct !{ !7 }
 !7 = distinct !{ !7, !6 }