!fixup address latest comments, thanks!

fhahn · fhahn · commit e26af8e878c6 · 2024-11-23T15:28:58.000Z
diff --git a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
@@ -7783,7 +7783,6 @@ DenseMap<const SCEV *, Value *> LoopVectorizationPlanner::executePlan(
     LoopVectorizeHints Hints(L, true, *ORE);
     Hints.setAlreadyVectorized();
   }
-
   TargetTransformInfo::UnrollingPreferences UP;
   TTI.getUnrollingPreferences(L, *PSE.getSE(), UP, ORE);
   if (!UP.UnrollVectorizedLoop || CanonicalIVStartValue)
@@ -8882,61 +8881,73 @@ static void addScalarResumePhis(VPRecipeBuilder &Builder, VPlan &Plan) {
 static SetVector<VPIRInstruction *> collectUsersInExitBlocks(
     Loop *OrigLoop, VPRecipeBuilder &Builder, VPlan &Plan,
     const MapVector<PHINode *, InductionDescriptor> &Inductions) {
+  auto *MiddleVPBB = Plan.getMiddleBlock();
   SetVector<VPIRInstruction *> ExitUsersToFix;
   for (VPIRBasicBlock *ExitVPBB : Plan.getExitBlocks()) {
-    BasicBlock *ExitBB = ExitVPBB->getIRBasicBlock();
     for (VPRecipeBase &R : *ExitVPBB) {
       auto *ExitIRI = dyn_cast<VPIRInstruction>(&R);
       if (!ExitIRI)
         continue;
       auto *ExitPhi = dyn_cast<PHINode>(&ExitIRI->getInstruction());
       if (!ExitPhi)
         break;
-      for (BasicBlock *ExitingBB : predecessors(ExitBB)) {
-        if (!OrigLoop->contains(ExitingBB))
-          continue;
-      Value *IncomingValue = ExitPhi->getIncomingValueForBlock(ExitingBB);
-      VPValue *V = Builder.getVPValueOrAddLiveIn(IncomingValue);
-      // Exit values for inductions are computed and updated outside of VPlan
-      // and independent of induction recipes.
-      // TODO: Compute induction exit values in VPlan.
-      if ((isa<VPWidenIntOrFpInductionRecipe>(V) &&
-           !cast<VPWidenIntOrFpInductionRecipe>(V)->getTruncInst()) ||
-          isa<VPWidenPointerInductionRecipe>(V) ||
-          (isa<Instruction>(IncomingValue) &&
-           OrigLoop->contains(cast<Instruction>(IncomingValue)) &&
-           any_of(IncomingValue->users(), [&Inductions](User *U) {
-             auto *P = dyn_cast<PHINode>(U);
-             return P && Inductions.contains(P);
-           })))
-        continue;
-      ExitUsersToFix.insert(ExitIRI);
-      ExitIRI->addOperand(V);
+      for (VPBlockBase *PredVPBB : ExitVPBB->getPredecessors()) {
+        BasicBlock *ExitingBB = OrigLoop->getLoopLatch();
+        if (PredVPBB != MiddleVPBB) {
+          SmallVector<BasicBlock *> ExitingBlocks;
+          OrigLoop->getExitingBlocks(ExitingBlocks);
+          assert(ExitingBlocks.size() == 2 && "only support 2 exiting blocks");
+          ExitingBB = ExitingBB == ExitingBlocks[0] ? ExitingBlocks[1]
+                                                    : ExitingBlocks[0];
+        }
+        Value *IncomingValue = ExitPhi->getIncomingValueForBlock(ExitingBB);
+        VPValue *V = Builder.getVPValueOrAddLiveIn(IncomingValue);
+        // Exit values for inductions are computed and updated outside of VPlan
+        // and independent of induction recipes.
+        // TODO: Compute induction exit values in VPlan.
+        if ((isa<VPWidenIntOrFpInductionRecipe>(V) &&
+             !cast<VPWidenIntOrFpInductionRecipe>(V)->getTruncInst()) ||
+            isa<VPWidenPointerInductionRecipe>(V) ||
+            (isa<Instruction>(IncomingValue) &&
+             OrigLoop->contains(cast<Instruction>(IncomingValue)) &&
+             any_of(IncomingValue->users(), [&Inductions](User *U) {
+               auto *P = dyn_cast<PHINode>(U);
+               return P && Inductions.contains(P);
+             }))) {
+          if (ExitVPBB->getSinglePredecessor() == MiddleVPBB)
+            continue;
+        }
+        ExitUsersToFix.insert(ExitIRI);
+        ExitIRI->addOperand(V);
       }
     }
   }
   return ExitUsersToFix;
 }
 
 // Add exit values to \p Plan. Extracts are added for each entry in \p
-// ExitUsersToFix if needed and their operands are updated.
-static void
+// ExitUsersToFix if needed and their operands are updated. Returns true if all
+// exit users can be handled, otherwise return false.
+static bool
 addUsersInExitBlocks(VPlan &Plan,
                      const SetVector<VPIRInstruction *> &ExitUsersToFix) {
   if (ExitUsersToFix.empty())
-    return;
+    return true;
+
+  auto *MiddleVPBB = Plan.getMiddleBlock();
+  VPBuilder B(MiddleVPBB, MiddleVPBB->getFirstNonPhi());
 
   // Introduce extract for exiting values and update the VPIRInstructions
   // modeling the corresponding LCSSA phis.
   for (VPIRInstruction *ExitIRI : ExitUsersToFix) {
-
     VPValue *V = ExitIRI->getOperand(0);
     // Pass live-in values used by exit phis directly through to their users in
     // the exit block.
     if (V->isLiveIn())
       continue;
 
-    // Currently only live-ins can be used by exit values from blocks not exiting via the vector latch through to the middle block.
+    // Currently only live-ins can be used by exit values from blocks not
+    // exiting via the vector latch through to the middle block.
     if (ExitIRI->getParent()->getSinglePredecessor() != MiddleVPBB)
       return false;
 
@@ -9218,31 +9229,19 @@ LoopVectorizationPlanner::tryToBuildVPlanWithVPRecipes(VFRange &Range) {
   RecipeBuilder.fixHeaderPhis();
 
   if (Legal->hasUncountableEarlyExit()) {
+    assert(Legal->getUncountableExitingBlocks().size() == 1 &&
+           "Only single uncountable exiting bock supported");
     VPlanTransforms::handleUncountableEarlyExit(
-        *Plan, *PSE.getSE(), OrigLoop, Legal->getUncountableExitingBlocks(),
+        *Plan, *PSE.getSE(), OrigLoop, Legal->getUncountableExitingBlocks()[0],
         RecipeBuilder);
   }
   addScalarResumePhis(RecipeBuilder, *Plan);
   SetVector<VPIRInstruction *> ExitUsersToFix = collectUsersInExitBlocks(
       OrigLoop, RecipeBuilder, *Plan, Legal->getInductionVars());
   addExitUsersForFirstOrderRecurrences(*Plan, ExitUsersToFix);
-  if (!addUsersInExitBlocks(*Plan, ExitUsersToFix)) 
-<<<<<<< HEAD
-
-  if (Legal->hasUncountableEarlyExit()) {
-    if (any_of(Plan->getExitBlocks(), [](VPIRBasicBlock *ExitBB) {
-          return any_of(*ExitBB, [](VPRecipeBase &R) {
-            auto VPIRI = cast<VPIRInstruction>(&R);
-            return VPIRI->getNumOperands() == 0 ||
-                   any_of(VPIRI->operands(),
-                          [](VPValue *Op) { return !Op->isLiveIn(); });
-          });
-        }))
-      return nullptr;
-  }
+  if (!addUsersInExitBlocks(*Plan, ExitUsersToFix))
+    return nullptr;
 
-=======
->>>>>>> origin/main
   // ---------------------------------------------------------------------------
   // Transform initial VPlan: Apply previously taken decisions, in order, to
   // bring the VPlan to its final state.
@@ -9998,11 +9997,10 @@ bool LoopVectorizePass::processLoop(Loop *L) {
   if (LVL.hasUncountableEarlyExit()) {
     if (!EnableEarlyExitVectorization) {
       reportVectorizationFailure("Auto-vectorization of loops with uncountable "
-                                 "early exit is not yet supported",
+                                 "early exit is not enabled",
                                  "Auto-vectorization of loops with uncountable "
-                                 "early exit is not yet supported",
-                                 "UncountableEarlyExitLoopsUnsupported", ORE,
-                                 L);
+                                 "early exit is no enabled",
+                                 "UncountableEarlyExitLoopsDisabled", ORE, L);
       return false;
     }
   }
diff --git a/llvm/lib/Transforms/Vectorize/VPlan.cpp b/llvm/lib/Transforms/Vectorize/VPlan.cpp
@@ -872,8 +872,10 @@ VPlanPtr VPlan::createInitialVPlan(Type *InductionTy,
 
   // Create SCEV and VPValue for the trip count.
   // We use the symbolic max backedge-taken-count, which is used when
-  // vectorizing loops with uncountable early exits
+  // vectorizing loops with uncountable early exits.
   const SCEV *BackedgeTakenCountSCEV = PSE.getSymbolicMaxBackedgeTakenCount();
+  assert(!isa<SCEVCouldNotCompute>(BackedgeTakenCountSCEV) &&
+         "Invalid loop count");
   ScalarEvolution &SE = *PSE.getSE();
   const SCEV *TripCount = SE.getTripCountFromExitCount(BackedgeTakenCountSCEV,
                                                        InductionTy, TheLoop);
@@ -907,16 +909,7 @@ VPlanPtr VPlan::createInitialVPlan(Type *InductionTy,
   // 2) If we require a scalar epilogue, there is no conditional branch as
   //    we unconditionally branch to the scalar preheader.  Do nothing.
   // 3) Otherwise, construct a runtime check.
-  BasicBlock *IRExitBlock = TheLoop->getUniqueExitBlock();
-  if (!IRExitBlock) {
-    // If there's no unique exit block (i.e. vectorizing with an uncountable
-    // early exit), use the block exiting from the latch. The other uncountable
-    // exit blocks will be added later.
-    auto *Term = cast<BranchInst>(TheLoop->getLoopLatch()->getTerminator());
-    IRExitBlock = TheLoop->contains(Term->getSuccessor(0))
-                      ? Term->getSuccessor(1)
-                      : Term->getSuccessor(0);
-  }
+  BasicBlock *IRExitBlock = TheLoop->getUniqueLatchExitBlock();
   auto *VPExitBlock = VPIRBasicBlock::fromBasicBlock(IRExitBlock);
   // The connection order corresponds to the operands of the conditional
   // branch.
diff --git a/llvm/lib/Transforms/Vectorize/VPlan.h b/llvm/lib/Transforms/Vectorize/VPlan.h
@@ -1232,6 +1232,8 @@ class VPInstruction : public VPRecipeWithIRFlags,
     // operand). Only generates scalar values (either for the first lane only or
     // for all lanes, depending on its uses).
     PtrAdd,
+    // Returns a scalar boolean value, which is true if any lane of its single
+    // operand is true.
     AnyOf,
   };
 
diff --git a/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp b/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp
@@ -68,8 +68,6 @@ bool VPRecipeBase::mayWriteToMemory() const {
     default:
       return true;
     }
-  case VPExpandSCEVSC:
-    return getParent()->getPlan()->getTripCount() == getVPSingleValue();
   case VPInterleaveSC:
     return cast<VPInterleaveRecipe>(this)->getNumStoreOperands() > 0;
   case VPWidenStoreEVLSC:
@@ -165,8 +163,6 @@ bool VPRecipeBase::mayHaveSideEffects() const {
   case VPScalarCastSC:
   case VPReverseVectorPointerSC:
     return false;
-  case VPExpandSCEVSC:
-    return getParent()->getPlan()->getTripCount() == getVPSingleValue();
   case VPInstructionSC:
     return mayWriteToMemory();
   case VPWidenCallSC: {
diff --git a/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp b/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp
@@ -509,12 +509,6 @@ void VPlanTransforms::removeDeadRecipes(VPlan &Plan) {
   ReversePostOrderTraversal<VPBlockDeepTraversalWrapper<VPBlockBase *>> RPOT(
       Plan.getEntry());
 
-  for (VPRecipeBase &R : make_early_inc_range(
-           reverse(*cast<VPBasicBlock>(Plan.getPreheader())))) {
-    if (isDeadRecipe(R))
-      R.eraseFromParent();
-  }
-
   for (VPBasicBlock *VPBB : reverse(VPBlockUtils::blocksOnly<VPBasicBlock>(RPOT))) {
     // The recipes in the block are processed in reverse order, to catch chains
     // of dead recipes.
@@ -1809,45 +1803,42 @@ void VPlanTransforms::createInterleaveGroups(
 
 void VPlanTransforms::handleUncountableEarlyExit(
     VPlan &Plan, ScalarEvolution &SE, Loop *OrigLoop,
-    ArrayRef<BasicBlock *> UncountableExitingBlocks,
-    VPRecipeBuilder &RecipeBuilder) {
+    BasicBlock *UncountableExitingBlock, VPRecipeBuilder &RecipeBuilder) {
   auto *LatchVPBB =
       cast<VPBasicBlock>(Plan.getVectorLoopRegion()->getExiting());
   VPBuilder Builder(LatchVPBB->getTerminator());
   auto *MiddleVPBB = Plan.getMiddleBlock();
   VPRegionBlock *LoopRegion = Plan.getVectorLoopRegion();
   VPValue *EarlyExitTaken = nullptr;
 
-  // Process all uncountable exiting blocks. For each exiting block, update the
-  // EarlyExitTaken, which tracks if any uncountable early exit has been taken.
-  // Also split the middle block and branch to the exit block for the early exit
-  // if it has been taken.
-  for (BasicBlock *Exiting : UncountableExitingBlocks) {
-    auto *ExitingTerm = cast<BranchInst>(Exiting->getTerminator());
-    BasicBlock *TrueSucc = ExitingTerm->getSuccessor(0);
-    BasicBlock *FalseSucc = ExitingTerm->getSuccessor(1);
-    VPIRBasicBlock *VPExitBlock;
-    if (OrigLoop->getUniqueExitBlock()) {
-      VPExitBlock = cast<VPIRBasicBlock>(MiddleVPBB->getSuccessors()[0]);
-    } else {
-      VPExitBlock = VPIRBasicBlock::fromBasicBlock(
-          !OrigLoop->contains(TrueSucc) ? TrueSucc : FalseSucc);
-    }
+  // Process the uncountable exiting block. Update EarlyExitTaken, which tracks
+  // if any uncountable early exit has been taken. Also split the middle block
+  // and branch to the exit block for the early exit if it has been taken.
+  auto *ExitingTerm =
+      cast<BranchInst>(UncountableExitingBlock->getTerminator());
+  BasicBlock *TrueSucc = ExitingTerm->getSuccessor(0);
+  BasicBlock *FalseSucc = ExitingTerm->getSuccessor(1);
+  VPIRBasicBlock *VPExitBlock;
+  if (OrigLoop->getUniqueExitBlock()) {
+    VPExitBlock = cast<VPIRBasicBlock>(MiddleVPBB->getSuccessors()[0]);
+  } else {
+    VPExitBlock = VPIRBasicBlock::fromBasicBlock(
+        !OrigLoop->contains(TrueSucc) ? TrueSucc : FalseSucc);
+  }
 
-    VPValue *M = RecipeBuilder.getBlockInMask(
-        OrigLoop->contains(TrueSucc) ? TrueSucc : FalseSucc);
-    auto *N = Builder.createNot(M);
-    EarlyExitTaken = Builder.createNaryOp(VPInstruction::AnyOf, {N});
+  VPValue *M = RecipeBuilder.getBlockInMask(
+      OrigLoop->contains(TrueSucc) ? TrueSucc : FalseSucc);
+  auto *N = Builder.createNot(M);
+  EarlyExitTaken = Builder.createNaryOp(VPInstruction::AnyOf, {N});
 
-    VPBasicBlock *NewMiddle = new VPBasicBlock("middle.split");
-    VPBlockUtils::disconnectBlocks(LoopRegion, MiddleVPBB);
-    VPBlockUtils::insertBlockAfter(NewMiddle, LoopRegion);
-    VPBlockUtils::connectBlocks(NewMiddle, VPExitBlock);
-    VPBlockUtils::connectBlocks(NewMiddle, MiddleVPBB);
+  VPBasicBlock *NewMiddle = new VPBasicBlock("middle.split");
+  VPBlockUtils::disconnectBlocks(LoopRegion, MiddleVPBB);
+  VPBlockUtils::insertBlockAfter(NewMiddle, LoopRegion);
+  VPBlockUtils::connectBlocks(NewMiddle, VPExitBlock);
+  VPBlockUtils::connectBlocks(NewMiddle, MiddleVPBB);
 
-    VPBuilder MiddleBuilder(NewMiddle);
-    MiddleBuilder.createNaryOp(VPInstruction::BranchOnCond, {EarlyExitTaken});
-  }
+  VPBuilder MiddleBuilder(NewMiddle);
+  MiddleBuilder.createNaryOp(VPInstruction::BranchOnCond, {EarlyExitTaken});
 
   // Replace the condition controlling the exit from the vector loop with one
   // exiting if either the original condition of the vector latch is true or any
diff --git a/llvm/lib/Transforms/Vectorize/VPlanTransforms.h b/llvm/lib/Transforms/Vectorize/VPlanTransforms.h
@@ -127,13 +127,13 @@ struct VPlanTransforms {
   /// Update \p Plan to account for uncountable exit blocks in \p
   /// UncountableExitingBlocks by
   ///  * updating the condition to exit the vector loop to include the early
-  ///  exit conditions
+  ///    exit conditions
   ///  * splitting the original middle block to branch to the early exit blocks
-  ///  if taken. Returns false if the transformation wasn't successful.
-  static void
-  handleUncountableEarlyExit(VPlan &Plan, ScalarEvolution &SE, Loop *OrigLoop,
-                             ArrayRef<BasicBlock *> UncountableExitingBlocks,
-                             VPRecipeBuilder &RecipeBuilder);
+  ///    if taken. Returns false if the transformation wasn't successful.
+  static void handleUncountableEarlyExit(VPlan &Plan, ScalarEvolution &SE,
+                                         Loop *OrigLoop,
+                                         BasicBlock *UncountableExitingBlock,
+                                         VPRecipeBuilder &RecipeBuilder);
 };
 
 } // namespace llvm
diff --git a/llvm/test/Transforms/LoopVectorize/early_exit_legality.ll b/llvm/test/Transforms/LoopVectorize/early_exit_legality.ll
@@ -1,6 +1,6 @@
 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 4
 ; REQUIRES: asserts
-; RUN: opt -S < %s -p loop-vectorize -debug-only=loop-vectorize -disable-output 2>&1 | FileCheck %s
+; RUN: opt -S < %s -p loop-vectorize -debug-only=loop-vectorize -enable-early-exit-vectorization -force-vector-width=4 -disable-output 2>&1 | FileCheck %s
 
 declare void @init_mem(ptr, i64);
 
diff --git a/llvm/test/Transforms/LoopVectorize/single_early_exit.ll b/llvm/test/Transforms/LoopVectorize/single_early_exit.ll
diff --git a/llvm/test/Transforms/LoopVectorize/uncountable-early-exit-vplan.ll b/llvm/test/Transforms/LoopVectorize/uncountable-early-exit-vplan.ll
diff --git a/llvm/test/Transforms/LoopVectorize/unsupported_early_exit.ll b/llvm/test/Transforms/LoopVectorize/unsupported_early_exit.ll
