[LV] Fix crash when optimized phi is feeding first-order recurrence.

[fhahn]

Summary

This patch fixes a crash when a phi with an optimized cast feeds a
first-order recurrence, as discussed in D113183.

There are 2 issues that this patch addresses:

1. The recipe for the optimized cast is not recorded.
2. fixHeaderPhis assumes the incoming def defines a single VPValue.

If the incoming value is an induction phi recipe, always use the first
defined value, which represents the phi itself.

[fhahn]

Removing of dead casts can be done directly in VPlan, so there's no need for special handling of induction recipes: D115112

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

VPRecipeBase *VPRecipeBuilder::tryToWidenMemory(Instruction *I,
 
   StoreInst *Store = cast<StoreInst>(I);
   return new VPWidenMemoryInstructionRecipe(*Store, Operands[1], Operands[0],
                                             Mask, Consecutive, Reverse);
 }
 
 VPWidenIntOrFpInductionRecipe *
 VPRecipeBuilder::tryToOptimizeInductionPHI(PHINode *Phi,
-                                           ArrayRef<VPValue *> Operands) const {
+                                           ArrayRef<VPValue *> Operands) {
   // Check if this is an integer or fp induction. If so, build the recipe that
   // produces its scalar and vector values.
   InductionDescriptor II = Legal->getInductionVars().lookup(Phi);
   if (II.getKind() == InductionDescriptor::IK_IntInduction ||
       II.getKind() == InductionDescriptor::IK_FpInduction) {
     assert(II.getStartValue() ==
            Phi->getIncomingValueForBlock(OrigLoop->getLoopPreheader()));
     const SmallVectorImpl<Instruction *> &Casts = II.getCastInsts();
-    return new VPWidenIntOrFpInductionRecipe(
-        Phi, Operands[0], Casts.empty() ? nullptr : Casts.front());
+
+    Instruction *CastI =
+        cast_or_null<Instruction>(Casts.empty() ? nullptr : Casts.front());
+    auto *R = new VPWidenIntOrFpInductionRecipe(Phi, Operands[0], CastI);
+    if (CastI)
+      setRecipe(CastI, R);
+    return R;
   }
 
   return nullptr;
 }
 
 VPWidenIntOrFpInductionRecipe *VPRecipeBuilder::tryToOptimizeInductionTruncate(
     TruncInst *I, ArrayRef<VPValue *> Operands, VFRange &Range,
     VPlan &Plan) const {

 }
 
 void VPRecipeBuilder::fixHeaderPhis() {
   BasicBlock *OrigLatch = OrigLoop->getLoopLatch();
   for (VPWidenPHIRecipe *R : PhisToFix) {
     auto *PN = cast<PHINode>(R->getUnderlyingValue());
     VPRecipeBase *IncR =
         getRecipe(cast<Instruction>(PN->getIncomingValueForBlock(OrigLatch)));
+    if (isa<VPWidenIntOrFpInductionRecipe>(IncR))
+      R->addOperand(IncR->getVPValue(0));
+    else
       R->addOperand(IncR->getVPSingleValue());
   }
 }
 
 VPBasicBlock *VPRecipeBuilder::handleReplication(
     Instruction *I, VFRange &Range, VPBasicBlock *VPBB,
     VPlanPtr &Plan) {
   bool IsUniform = LoopVectorizationPlanner::getDecisionAndClampRange(
       [&](ElementCount VF) { return CM.isUniformAfterVectorization(I, VF); },

llvm/lib/Transforms/Vectorize/VPRecipeBuilder.h

class VPRecipeBuilder {
   /// Range.Start and potentially masked. Such instructions are handled by a
   /// recipe that takes an additional VPInstruction for the mask.
   VPRecipeBase *tryToWidenMemory(Instruction *I, ArrayRef<VPValue *> Operands,
                                  VFRange &Range, VPlanPtr &Plan);
 
   /// Check if an induction recipe should be constructed for \I. If so build and
   /// return it. If not, return null.
   VPWidenIntOrFpInductionRecipe *
-  tryToOptimizeInductionPHI(PHINode *Phi, ArrayRef<VPValue *> Operands) const;
+  tryToOptimizeInductionPHI(PHINode *Phi, ArrayRef<VPValue *> Operands);
 
   /// Optimize the special case where the operand of \p I is a constant integer
   /// induction variable.
   VPWidenIntOrFpInductionRecipe *
   tryToOptimizeInductionTruncate(TruncInst *I, ArrayRef<VPValue *> Operands,
                                  VFRange &Range, VPlan &Plan) const;
 
   /// Handle non-loop phi nodes. Return a VPValue, if all incoming values match

llvm/test/Transforms/LoopVectorize/induction.ll

loop:
   %add = add i32 %iv.trunc, %mul
   store i32 %add, i32* %dst.gep
   %exitcond = icmp eq i32 %trunc.iv.next, 100
   br i1 %exitcond, label %exit, label %loop
 
 exit:
   ret void
 }
+
+; Test case where %iv.2.ext and %iv.2.conv become redundant due to the SCEV
+; predicates generated for the vector loop. They should be removed in the
+; vector loop.
+define void @test_optimized_cast_induction_feeding_first_order_recurrence(i64 %n, i32 %step, i32* %ptr) {
+; CHECK-LABEL: @test_optimized_cast_induction_feeding_first_order_recurrence(
+; CHECK-LABEL: vector.body:
+; CHECK-NEXT:    [[MAIN_IV:%.+]] = phi i64 [ 0, %vector.ph ], [ [[MAIN_IV_NEXT:%.+]], %vector.body ]
+; CHECK-NEXT:    [[VEC_RECUR:%.+]] = phi <2 x i32> [ <i32 poison, i32 0>, %vector.ph ], [ [[VEC_IV:%.+]], %vector.body ]
+; CHECK-NEXT:    [[VEC_IV]] = phi <2 x i32> [ %induction, %vector.ph ], [ [[VEC_IV_NEXT:%.+]], %vector.body ]
+; CHECK-NEXT:    [[MAIN_IV_0:%.+]] = add i64 [[MAIN_IV]], 0
+; CHECK-NEXT:    [[RECUR_SHUFFLE:%.+]] = shufflevector <2 x i32> [[VEC_RECUR]], <2 x i32> [[VEC_IV]], <2 x i32> <i32 1, i32 2>
+; CHECK-NEXT:    [[GEP0:%.+]] = getelementptr inbounds i32, i32* %ptr, i64 [[MAIN_IV_0]]
+; CHECK-NEXT:    [[GEP1:%.+]] = getelementptr inbounds i32, i32* [[GEP0]], i32 0
+; CHECK-NEXT:    [[GEP_CAST:%.+]] = bitcast i32* [[GEP1]] to <2 x i32>*
+; CHECK-NEXT:    store <2 x i32> [[RECUR_SHUFFLE]], <2 x i32>* [[GEP_CAST]], align 4
+; CHECK-NEXT:    [[MAIN_IV_NEXT]] = add nuw i64 [[MAIN_IV]], 2
+; CHECK-NEXT:    [[VEC_IV_NEXT]] = add <2 x i32> [[VEC_IV]],
+;
+entry:
+  br label %loop
+
+loop:
+  %for = phi i32 [ 0, %entry ], [ %iv.2.conv, %loop ]
+  %iv.1 = phi i64 [ 0, %entry ], [ %iv.1.next, %loop ]
+  %iv.2 = phi i32 [ 0, %entry ], [ %iv.2.next, %loop ]
+  %iv.2.ext = shl i32 %iv.2, 24
+  %iv.2.conv = ashr exact i32 %iv.2.ext, 24
+  %gep = getelementptr inbounds i32, i32* %ptr, i64 %iv.1
+  store i32 %for, i32* %gep, align 4
+  %iv.2.next = add nsw i32 %iv.2.conv, %step
+  %iv.1.next = add nuw nsw i64 %iv.1, 1
+  %exitcond = icmp eq i64 %iv.1.next, %n
+  br i1 %exitcond, label %exit, label %loop
+
+exit:
+  ret void
+}