[LV] Fix crash when target instruction for sinking is dead.

[fhahn]

Summary

This patch fixes a crash when the target instruction for sinking is
dead. In that case, no recipe is created and trying to get the recipe
for it results in a crash. To ensure all sink targets are alive, find &
use the first previous alive instruction.

Note that the case where the sink source is dead is already handled.

Found by
https://bugs.chromium.org/p/oss-fuzz/issues/detail?id=35320

[Ayal]

Ahh, the culprit here is a chain of FOR phi-users all sinking after a Previous - the latter should never be dead, but any member of the chain may by dead.
In the test case: %rec.2 is the FOR phi, %rec.2.prev is Previous, and {%cmp, %C, %B2} is the chain of users, where %B2 is dead as it feeds the loop's conditional branch only.

An alternative fix is to have RecurrenceDescriptor::isFirstOrderRecurrence() sink the chain in reverse, always after the non-dead Previous; as you originally suggested in D84951?

[fhahn]

In D104603#2830492, @Ayal wrote:

Ahh, the culprit here is a chain of FOR phi-users all sinking after a Previous - the latter should never be dead, but any member of the chain may by dead.
In the test case: %rec.2 is the FOR phi, %rec.2.prev is Previous, and {%cmp, %C, %B2} is the chain of users, where %B2 is dead as it feeds the loop's conditional branch only.

Yes exactly! Let me improve the naming of the variables a bit.

An alternative fix is to have RecurrenceDescriptor::isFirstOrderRecurrence() sink the chain in reverse, always after the non-dead Previous; as you originally suggested in D84951?

I think another alternative would be to handle the check for dead instructions directly at the place where we do sinking. I updated the code to do that. Please let me know if you prefer the alternative.

I think chaining the instructions in the map makes sense and the handling of dead instructions before VPlan creation may go away in the medium term anyways?

[fhahn]

Adjust FOR phi variable name in test.

[Ayal]

In D104603#2831278, @fhahn wrote:

In D104603#2830492, @Ayal wrote:

Ahh, the culprit here is a chain of FOR phi-users all sinking after a Previous - the latter should never be dead, but any member of the chain may by dead.
In the test case: %rec.2 is the FOR phi, %rec.2.prev is Previous, and {%cmp, %C, %B2} is the chain of users, where %B2 is dead as it feeds the loop's conditional branch only.

Yes exactly! Let me improve the naming of the variables a bit.

Good! Can also call out explicitly the %dead value.

An alternative fix is to have RecurrenceDescriptor::isFirstOrderRecurrence() sink the chain in reverse, always after the non-dead Previous; as you originally suggested in D84951?

I think another alternative would be to handle the check for dead instructions directly at the place where we do sinking. I updated the code to do that. Please let me know if you prefer the alternative.

I think chaining the instructions in the map makes sense and the handling of dead instructions before VPlan creation may go away in the medium term anyways?

Agreed that in the medium term, dead instructions may be handled differently in/before VPlan/Legal, so would be good to have a simple fix for now.

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

9041

It may be better to handle all updates of SinkAfter due to DeadInstructions in one place, i.e., erasing dead sinks and bumping dead targets, either both here or both below inside buildVPlanWithVPRecipes(Range).

9224

Traversing the dead backwards here should also be ok, but raises additional thoughts.
If we reach FirstInst == SinkTarget then no sinking is required, which should not happen, so best have an assert instead? May want to also assert against falling off the first recipe of the block.

Additional alternatives, on top of reversing SinkAfter chains, include (mainly for completion):

1. Letting Legal handle DeadInstructions, so that SinkTarget will include only live sinks to be moved after live targets. But there's a dependence on FoldTail wanting to "reuse" Primary Induction.
2. Having the first VPlan wrap Dead Instructions with a (recorded) Dead Recipe, to be eliminated by a VPlan-to-VPlan DCE.
3. Set LastLiveSecond = (Entry.second is live ? Entry.second : LastLiveSecond), and always sink SinkTarget after LastLiveSecond. This relies on iteration order over SinkAfter being {U_1,Previous},{U_2,U_1},...,{U_k,U_{k-1}} whenever some target U_i is dead, where Previous must be live.

Can also handle dead sinks here, as mentioned above, e.g., by continue. Less efficient, but sizes of SinkAfter and DeadInstructions should be very small if positive.

Can also update the sinkAfter pair whose target was revived, caching the result for next Range.

[fhahn]

Moved back to location next to deleting dead sink sources. Also added assertions as suggested.

In D104603#2842826, @Ayal wrote:

In D104603#2831278, @fhahn wrote:

In D104603#2830492, @Ayal wrote:

Ahh, the culprit here is a chain of FOR phi-users all sinking after a Previous - the latter should never be dead, but any member of the chain may by dead.
In the test case: %rec.2 is the FOR phi, %rec.2.prev is Previous, and {%cmp, %C, %B2} is the chain of users, where %B2 is dead as it feeds the loop's conditional branch only.

Yes exactly! Let me improve the naming of the variables a bit.

Good! Can also call out explicitly the %dead value.

Done, also added a comment to the test.

[fhahn]

Rename *Cur* to *SinkTarget*.

[fhahn]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

9224

Traversing the dead backwards here should also be ok, but raises additional thoughts.
If we reach FirstInst == SinkTarget then no sinking is required, which should not happen, so best have an assert instead? May want to also assert against falling off the first recipe of the block.

I added 2 assertions above:

• SinkTarget != FirstInst, should guard against failing of the first recipe in the block
• SinkTarget != SinkSource, should guard against cases not requiring any sinking?

WDYT?

[Ayal]

Thanks!

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

9052

a alive >> a live

9054

This second assert(SinkTarget != P.first) should be placed right after bumping SinkTarget to its prev node.
(We know P.first ain't dead - these were taken care of above; so if SinkTarget get set to P.first, we will immediately exit the loop, before reaching this assert here. So can also place the assert right after the loop.)
(The first assert is placed well, guarding against taking FirstInst->getPrevNode().)

9224

+1, just a remark above about where to place the second assert.

[fhahn]

Thanks Ayal!

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

9054

Thanks, I moved the second assert in the committed version.

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

void LoopVectorizationPlanner::buildVPlansWithVPRecipes(ElementCount MinVF,
   // TODO: We only need to drop assumes in blocks that get flattend. If the
   // control flow is preserved, we should keep them.
   auto &ConditionalAssumes = Legal->getConditionalAssumes();
   DeadInstructions.insert(ConditionalAssumes.begin(), ConditionalAssumes.end());
 
   MapVector<Instruction *, Instruction *> &SinkAfter = Legal->getSinkAfter();
   // Dead instructions do not need sinking. Remove them from SinkAfter.
   for (Instruction *I : DeadInstructions)
     SinkAfter.erase(I);
 
+  // Cannot sink instructions after dead instructions (there won't be any
+  // recipes for them). Instead, find the first non-dead previous instruction.
+  for (auto &P : Legal->getSinkAfter()) {
+    Instruction *SinkTarget = P.second;
+    Instruction *FirstInst = &*SinkTarget->getParent()->begin();
+    (void)FirstInst;
+    while (DeadInstructions.contains(SinkTarget)) {
+      assert(
+          SinkTarget != FirstInst &&
+          "Must find a live instruction (at least the one feeding the "
+          "first-order recurrence PHI) before reaching beginning of the block");
+      SinkTarget = SinkTarget->getPrevNode();
+      assert(SinkTarget != P.first &&
+             "sink source equals target, no sinking required");
+    }
+    P.second = SinkTarget;
+  }
+
   auto MaxVFPlusOne = MaxVF.getWithIncrement(1);
   for (ElementCount VF = MinVF; ElementCount::isKnownLT(VF, MaxVFPlusOne);) {
     VFRange SubRange = {VF, MaxVFPlusOne};
     VPlans.push_back(
         buildVPlanWithVPRecipes(SubRange, DeadInstructions, SinkAfter));
     VF = SubRange.End;
   }
 }
VPlanPtr LoopVectorizationPlanner::buildVPlanWithVPRecipes(
 
   // Apply Sink-After legal constraints.
   for (auto &Entry : SinkAfter) {
     VPRecipeBase *Sink = RecipeBuilder.getRecipe(Entry.first);
     VPRecipeBase *Target = RecipeBuilder.getRecipe(Entry.second);
 
     auto GetReplicateRegion = [](VPRecipeBase *R) -> VPRegionBlock * {
       auto *Region =
           dyn_cast_or_null<VPRegionBlock>(R->getParent()->getParent());
       if (Region && Region->isReplicator()) {
         assert(Region->getNumSuccessors() == 1 &&
                Region->getNumPredecessors() == 1 && "Expected SESE region!");
         assert(R->getParent()->size() == 1 &&
                "A recipe in an original replicator region must be the only "
                "recipe in its block");
         return Region;
       }

llvm/test/Transforms/LoopVectorize/first-order-recurrence.ll

bb:
   %tmp7 = udiv i32 219220132, %tmp3
   store i32 %tmp3, i32* %g, align 4
   %tmp8 = add nsw i32 %tmp3, -1
   %iv.next = add nsw i32 %iv, 1
   %tmp9 = icmp slt i32 %tmp3, 2
   br i1 %tmp9, label %bb1, label %bb2, !prof !2
 }
 
+; %vec.dead will be marked as dead instruction in the vector loop and no recipe
+; will be created for it. Make sure a valid sink target is used.
+define void @sink_after_dead_inst(i32* %A.ptr) {
+; CHECK-LABEL: @sink_after_dead_inst
+; CHECK-LABEL: vector.body:
+; CHECK-NEXT:    [[INDEX:%.*]] = phi i32 [ 0, %vector.ph ], [ [[INDEX_NEXT]], %vector.body ]
+; CHECK-NEXT:    [[OFFSET_IDX:%.*]] = zext i32 [[INDEX]] to i64
+; CHECK-NEXT:    [[SEXT:%.*]] = shl i64 [[OFFSET_IDX]], 48
+; CHECK-NEXT:    [[SHIFT:%.*]] = ashr exact i64 [[SEXT]], 48
+; CHECK-NEXT:    [[GEP:%.*]] = getelementptr i32, i32* %A.ptr, i64 [[SHIFT]]
+; CHECK-NEXT:    [[CAST:%.*]] = bitcast i32* [[GEP]]  to <4 x i32>*
+; CHECK-NEXT:    store <4 x i32> zeroinitializer, <4 x i32>* [[CAST]], align 4
+; CHECK-NEXT:    [[INDEX_NEXT:%.*]] = add nuw i32 [[INDEX]], 4
+; CHECK-NEXT:    [[EC:%.*]] = icmp eq i32 [[INDEX_NEXT]], 16
+; CHECK-NEXT:    br i1 [[EC]], label %middle.block, label %vector.body
+
+entry:
+  br label %loop
+
+loop:
+  %iv = phi i16 [ 0, %entry ], [ %iv.next, %loop ]
+  %for = phi i32 [ 0, %entry ], [ %for.prev, %loop ]
+  %cmp = icmp eq i32 %for, 15
+  %C = icmp eq i1 %cmp, true
+  %vec.dead = and i1 %C, 1
+  %iv.next = add i16 %iv, 1
+  %B1 = or i16 %iv.next, %iv.next
+  %B3 = and i1 %cmp, %C
+  %for.prev = zext i16 %B1 to i32
+
+  %ext = zext i1 %B3 to i32
+  %A.gep = getelementptr i32, i32* %A.ptr, i16 %iv
+  store i32 0, i32* %A.gep
+  br i1 %vec.dead, label %for.end, label %loop
+
+for.end:
+  ret void
+}
+
 !2 = !{!"branch_weights", i32 1, i32 1}