[VPlan] Make sure optimizeInductions removes wide ind from scalar plan.

[fhahn]

Summary

In some cases, there may be widened users of inductions even though the
plan includes the scalar VF. In those cases, make sure we still replace
the VPWidenIntOrFpInductionRecipe with scalar steps, as otherwise we may
try to execute a VPWidenIntOrFpInductionRecipe with a scalar VF.

Alternatively the patch could also split the range if needed.

This fixes a crash exposed by D123720.

[bmahjour]

llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp

387–390

suggestion:

(none_of(IV->users(), [IV](VPUser *U) { return U->usesScalars(IV); })

[bmahjour]

could you please add the crashing IR from https://reviews.llvm.org/D123720 as an additional test? It still seems to crash with this patch!

[Ayal]

llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp

384

How about AllUseScalars, or ScalarVF? "NeedsScalarSteps" combines the existence of a true usesScalars user together with a scalar VF.

Worth clarifying that usesScalars() assumes a VF>1 vector is given, as when VF=1 all users are provided a scalar and use it. I.e., "usesScalarsOfAVector()".

388

Worth reversing the condition, to early-continue if (!(IV && (NeedsScalarSteps || any_of(usesScalars(IV))?
(BTW, none_of() seems clearer than all_of(!)).

406

Worth reversing the condition, to early-continue if (!(NeedsScalarSteps || usesScalars(IV))?

llvm/test/Transforms/LoopVectorize/interleave-and-scalarize-only.ll

172

Explain what this test is designed to check? (Worth pre-committing it to demonstrate failure?)

[fhahn]

Address latest comments, thanks!

In D128755#3617144, @bmahjour wrote:

could you please add the crashing IR from https://reviews.llvm.org/D123720 as an additional test? It still seems to crash with this patch!

Is it still crashing with this patch applied to current main? It needs b18141a8f29f638be0602aa6ffbfc2d434b0b74a as well.

[fhahn]

llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp

384

How about AllUseScalars, or ScalarVF? "NeedsScalarSteps" combines the existence of a true usesScalars user together with a scalar VF.

Good point, I updated the variable name to ScalarVF.

Worth clarifying that usesScalars() assumes a VF>1 vector is given, as when VF=1 all users are provided a scalar and use it. I.e., "usesScalarsOfAVector()".

Sounds good, I can update that separately.

387–390

Thanks, I wasn't aware of none_of. Updated!

388

Worth reversing the condition, to early-continue if (!(IV && (NeedsScalarSteps || any_of(usesScalars(IV))?
(BTW, none_of() seems clearer than all_of(!)).

I moved the ScarlarVF check first, thanks! I also updated the check to use none_of and also moved out the !IV check to a separate if which hopefully should make the code a bit easier to follow.

llvm/test/Transforms/LoopVectorize/interleave-and-scalarize-only.ll

172

Thanks, I added an explanation. Unfortunately pre-committing it would mean we have to XFAIL the whole file, as it causes a crash with assertions :(

[bmahjour]

In D128755#3618609, @fhahn wrote:

Address latest comments, thanks!

In D128755#3617144, @bmahjour wrote:

could you please add the crashing IR from https://reviews.llvm.org/D123720 as an additional test? It still seems to crash with this patch!

Is it still crashing with this patch applied to current main? It needs b18141a8f29f638be0602aa6ffbfc2d434b0b74a as well.

I was indeed missing b18141a8f29f638be0602aa6ffbfc2d434b0b74a. I've verified that the crash is gone with this patch on top of latest trunk. Thank you!

[Ayal]

This fix is fine, thanks, adding a minor nit.

llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp

384

Ah, perhaps HasVectorVF = !Plan

388

Worth reversing the condition, to early-continue if (!(IV && (NeedsScalarSteps || any_of(usesScalars(IV))?
(BTW, none_of() seems clearer than all_of(!)).

I moved the ScarlarVF check first, thanks! I also updated the check to use none_of and also moved out the !IV check to a separate if which hopefully should make the code a bit easier to follow.

Thanks!
nit: seems logical to combine HasScalarVF || any_of(usesScalars(), because a scalar VF implies all use scalars. and then reverse it, w/o applying De Morgan, here and below. But if you prefer the current form is also ok.

[fhahn]

llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp

384

I updated the name to use HasOnlyVectorVFs, thanks!

llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp

for (VPRecipeBase &R : make_early_inc_range(reverse(*VPBB))) {
       R.eraseFromParent();
     }
   }
 }
 
 void VPlanTransforms::optimizeInductions(VPlan &Plan, ScalarEvolution &SE) {
   SmallVector<VPRecipeBase *> ToRemove;
   VPBasicBlock *HeaderVPBB = Plan.getVectorLoopRegion()->getEntryBasicBlock();
+  bool HasOnlyVectorVFs = !Plan.hasVF(ElementCount::getFixed(1));
   for (VPRecipeBase &Phi : HeaderVPBB->phis()) {
     auto *IV = dyn_cast<VPWidenIntOrFpInductionRecipe>(&Phi);
-    if (!IV ||
-        all_of(IV->users(), [IV](VPUser *U) { return !U->usesScalars(IV); }))
+    if (!IV)
+      continue;
+    if (HasOnlyVectorVFs &&
+        none_of(IV->users(), [IV](VPUser *U) { return U->usesScalars(IV); }))
       continue;
 
     const InductionDescriptor &ID = IV->getInductionDescriptor();
     VPValue *Step =
         vputils::getOrCreateVPValueForSCEVExpr(Plan, ID.getStep(), SE);
     Instruction *TruncI = IV->getTruncInst();
     VPScalarIVStepsRecipe *Steps = new VPScalarIVStepsRecipe(
         IV->getPHINode()->getType(), ID, Plan.getCanonicalIV(),
         IV->getStartValue(), Step, TruncI ? TruncI->getType() : nullptr);
     HeaderVPBB->insert(Steps, HeaderVPBB->getFirstNonPhi());
 
     // Update scalar users of IV to use Step instead. Use SetVector to ensure
     // the list of users doesn't contain duplicates.
     SetVector<VPUser *> Users(IV->user_begin(), IV->user_end());
     for (VPUser *U : Users) {
-      if (!U->usesScalars(IV))
+      if (HasOnlyVectorVFs && !U->usesScalars(IV))
         continue;
       for (unsigned I = 0, E = U->getNumOperands(); I != E; I++) {
         if (U->getOperand(I) != IV)
           continue;
         U->setOperand(I, Steps);
       }
     }
   }

llvm/test/Transforms/LoopVectorize/interleave-and-scalarize-only.ll

 loop.latch:
   %iv.next = add nsw i64 %iv, 1
   %ec = icmp eq i64 %iv.next, 1000
   br i1 %ec, label %exit, label %loop.header
 
 exit:
   ret void
 }
+
+; Make sure the widened induction gets replaced by scalar-steps for plans
+; including the scalar VF, if it is used in first-order recurrences.
+
+; DBG-LABEL: 'first_order_recurrence_using_induction'
+; DBG:      VPlan 'Initial VPlan for VF={1},UF>=1' {
+; DBG-NEXT: Live-in vp<%1> = vector-trip-count
+; DBG-EMPTY:
+; DBG-NEXT: vector.ph:
+; DBG-NEXT: Successor(s): vector loop
+; DBG-EMPTY:
+; DBG-NEXT: <x1> vector loop: {
+; DBG-NEXT:   vector.body:
+; DBG-NEXT:     EMIT vp<%2> = CANONICAL-INDUCTION
+; DBG-NEXT:     FIRST-ORDER-RECURRENCE-PHI ir<%for> = phi ir<0>, vp<%4>
+; DBG-NEXT:     vp<%4>    = SCALAR-STEPS vp<%2>, ir<0>, ir<1>
+; DBG-NEXT:     EMIT vp<%5> = first-order splice ir<%for> vp<%4>
+; DBG-NEXT:     CLONE store vp<%5>, ir<%dst>
+; DBG-NEXT:     EMIT vp<%7> = VF * UF +(nuw)  vp<%2>
+; DBG-NEXT:     EMIT branch-on-count  vp<%7> vp<%1>
+; DBG-NEXT:   No successors
+; DBG-NEXT: }
+; DBG-NEXT: Successor(s): middle.block
+; DBG-EMPTY:
+; DBG-NEXT: middle.block:
+; DBG-NEXT: No successors
+; DBG-NEXT: }
+
+define void @first_order_recurrence_using_induction(i32 %n, ptr %dst) {
+; CHECK-LABEL: @first_order_recurrence_using_induction(
+; CHECK:       vector.body:
+; CHECK-NEXT:    [[INDEX:%.*]] = phi i64 [ 0, %vector.ph ], [ [[INDEX_NEXT:%.*]], %vector.body ]
+; CHECK-NEXT:    [[VECTOR_RECUR:%.*]] = phi i32 [ 0, %vector.ph ], [ [[INDUCTION1:%.*]], %vector.body ]
+; CHECK-NEXT:    [[TMP3:%.*]] = trunc i64 [[INDEX]] to i32
+; CHECK-NEXT:    [[INDUCTION:%.*]] = add i32 [[TMP3]], 0
+; CHECK-NEXT:    [[INDUCTION1]] = add i32 [[TMP3]], 1
+; CHECK-NEXT:    store i32 [[VECTOR_RECUR]], ptr [[DST:%.*]], align 4
+; CHECK-NEXT:    store i32 [[INDUCTION]], ptr [[DST]], align 4
+; CHECK-NEXT:    [[INDEX_NEXT]] = add nuw i64 [[INDEX]], 2
+; CHECK-NEXT:    [[TMP4:%.*]] = icmp eq i64 [[INDEX_NEXT]], %n.vec
+; CHECK-NEXT:    br i1 [[TMP4]], label %middle.block, label %vector.body
+; CHECK:       middle.block:
+;
+entry:
+  br label %loop
+
+loop:
+  %iv = phi i64 [ 0, %entry ],[ %iv.next, %loop ]
+  %for = phi i32 [ 0, %entry ], [ %iv.trunc, %loop ]
+  %iv.trunc = trunc i64 %iv to i32
+  store i32 %for, ptr %dst
+  %iv.next = add nuw nsw i64 %iv, 1
+  %iv.next.trunc = trunc i64 %iv.next to i32
+  %ec = icmp slt i32 %iv.next.trunc, %n
+  br i1 %ec, label %loop, label %exit
+
+exit:
+  ret void
+}