!fixup introduce and use VPIRInstruction::extractLastLaneOfOperand.

Author: fhahn

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

@@ -9095,16 +9095,11 @@ addUsersInExitBlocks(VPlan &Plan,
   // Introduce extract for exiting values and update the VPIRInstructions
   // modeling the corresponding LCSSA phis.
   for (VPIRInstruction *ExitIRI : ExitUsersToFix) {
-    for (const auto &[Idx, Op] : enumerate(ExitIRI->operands())) {
-      assert(ExitIRI->getParent()->getSinglePredecessor() == MiddleVPBB &&
-             "exit values from early exits must be fixed when branch to "
-             "early-exit is added");
-      LLVMContext &Ctx = ExitIRI->getInstruction().getContext();
-      VPValue *Ext = B.createNaryOp(VPInstruction::ExtractFromEnd,
-                                    {Op, Plan.getOrAddLiveIn(ConstantInt::get(
-                                             IntegerType::get(Ctx, 32), 1))});
-      ExitIRI->setOperand(Idx, Ext);
-    }
+    assert(ExitIRI->getNumOperands() == 1 &&
+           ExitIRI->getParent()->getSinglePredecessor() == MiddleVPBB &&
+           "exit values from early exits must be fixed when branch to "
+           "early-exit is added");
+    ExitIRI->extractLastLaneOfOperand(B);
   }
 }
 

llvm/lib/Transforms/Vectorize/VPlan.h

@@ -60,6 +60,7 @@ class RecurrenceDescriptor;
 class SCEV;
 class Type;
 class VPBasicBlock;
+class VPBuilder;
 class VPRegionBlock;
 class VPlan;
 class VPReplicateRecipe;
@@ -1428,6 +1429,11 @@ class VPIRInstruction : public VPRecipeBase {
            "Op must be an operand of the recipe");
     return true;
   }
+
+  /// Update the recipes single operand to the last lane of the operand using \p
+  /// Builder. Must only be used for single operand VPIRInstructions wrapping a
+  /// PHINode.
+  void extractLastLaneOfOperand(VPBuilder &Builder);
 };
 
 /// VPWidenRecipe is a recipe for producing a widened instruction using the

llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp

@@ -11,6 +11,7 @@
 ///
 //===----------------------------------------------------------------------===//
 
+#include "LoopVectorizationPlanner.h"
 #include "VPlan.h"
 #include "VPlanAnalysis.h"
 #include "VPlanPatternMatch.h"
@@ -937,6 +938,22 @@ InstructionCost VPIRInstruction::computeCost(ElementCount VF,
   return 0;
 }
 
+void VPIRInstruction::extractLastLaneOfOperand(VPBuilder &Builder) {
+  assert(isa<PHINode>(getInstruction()) &&
+         "can only add exiting operands to phi nodes");
+  assert(getNumOperands() == 1 && "must have a single operand");
+  VPValue *Exiting = getOperand(0);
+  if (!Exiting->isLiveIn()) {
+    LLVMContext &Ctx = getInstruction().getContext();
+    auto &Plan = *getParent()->getPlan();
+    Exiting = Builder.createNaryOp(
+        VPInstruction::ExtractFromEnd,
+        {Exiting,
+         Plan.getOrAddLiveIn(ConstantInt::get(IntegerType::get(Ctx, 32), 1))});
+  }
+  setOperand(0, Exiting);
+}
+
 #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
 void VPIRInstruction::print(raw_ostream &O, const Twine &Indent,
                             VPSlotTracker &SlotTracker) const {

llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp

@@ -2114,7 +2114,7 @@ bool VPlanTransforms::handleUncountableEarlyExit(
     VPValue *IncomingFromEarlyExit = RecipeBuilder.getVPValueOrAddLiveIn(
         ExitPhi->getIncomingValueForBlock(UncountableExitingBlock));
     // The incoming value from the early exit must be a live-in for now.
-    if (!IncomingFromm EarlyExit->isLiveIn())
+    if (!IncomingFromEarlyExit->isLiveIn())
       return false;
 
     if (OrigLoop->getUniqueExitBlock()) {
@@ -2123,14 +2123,8 @@ bool VPlanTransforms::handleUncountableEarlyExit(
       // which is coming from the original latch.
       VPValue *IncomingFromLatch = RecipeBuilder.getVPValueOrAddLiveIn(
           ExitPhi->getIncomingValueForBlock(OrigLoop->getLoopLatch()));
-      if (!IncomingFromLatch->isLiveIn()) {
-        LLVMContext &Ctx = ExitIRI->getInstruction().getContext();
-        IncomingFromLatch = MiddleBuilder.createNaryOp(
-            VPInstruction::ExtractFromEnd,
-            {IncomingFromLatch, Plan.getOrAddLiveIn(ConstantInt::get(
-                                    IntegerType::get(Ctx, 32), 1))});
-      }
       ExitIRI->addOperand(IncomingFromLatch);
+      ExitIRI->extractLastLaneOfOperand(MiddleBuilder);
     }
     // Add the incoming value from the early exit.
     ExitIRI->addOperand(IncomingFromEarlyExit);