[llvm] [LV] Transform to handle exits in the scalar loop (PR #148626)

[Gaëtan Bossu via llvm-commits]

================
@@ -3610,6 +3616,147 @@ void VPlanTransforms::handleUncountableEarlyExit(VPBasicBlock *EarlyExitingVPBB,
   LatchExitingBranch->eraseFromParent();
 }
 
+bool VPlanTransforms::handleUncountableExitsInScalarLoop(VPlan &Plan) {
+  assert(!Plan.hasScalarVFOnly() &&
+         "Cannot transform uncountable exits in scalar loop");
+
+  // We can abandon a vplan entirely if we return false here, so we shouldn't
+  // crash if some earlier assumptions on scalar IR don't hold for the vplan
+  // version of the loop.
+  VPCanonicalIVPHIRecipe *IV = Plan.getVectorLoopRegion()->getCanonicalIV();
+  VPInstruction *IVUpdate = dyn_cast<VPInstruction>(IV->getBackedgeValue());
+  if (!IVUpdate)
+    return false;
+
+  SmallVector<VPRecipeBase *, 2> GEPs;
+  SmallVector<VPRecipeBase *, 8> ConditionRecipes;
+
+  std::optional<VPValue *> Cond =
+      vputils::getRecipesForUncountableExit(Plan, ConditionRecipes, GEPs);
+  if (!Cond)
+    return false;
+
+  // Check GEPs to see if we can link them to the canonical IV.
+  using namespace llvm::VPlanPatternMatch;
+  for (auto *GEP : GEPs)
+    if (!match(GEP,
+               m_GetElementPtr(m_LiveIn(),
+                               m_ScalarIVSteps(m_Specific(IV), m_SpecificInt(1),
+                                               m_Specific(&Plan.getVF())))))
+      return false;
+
+  // Clone the condition recipes into the preheader
+  SmallDenseMap<VPRecipeBase *, VPRecipeBase *, 8> CloneMap;
+  VPBasicBlock *VectorPH = Plan.getVectorPreheader();
+  for (VPRecipeBase *R : reverse(ConditionRecipes)) {
+    VPRecipeBase *Clone = R->clone();
+    VectorPH->appendRecipe(Clone);
+    CloneMap[R] = Clone;
+  }
+
+  // Remap the cloned recipes to use the corresponding operands.
+  for (VPRecipeBase *R : ConditionRecipes) {
+    auto *Clone = CloneMap.at(R);
+    for (unsigned I = 0; I < R->getNumOperands(); ++I)
+      if (VPRecipeBase *OpR =
+              CloneMap.lookup(R->getOperand(I)->getDefiningRecipe()))
+        Clone->setOperand(I, OpR->getVPSingleValue());
+  }
+
+  // Adjust preheader GEPs to match the value they would have for the first
+  // iteration of the vector body.
+  for (auto *GEP : GEPs)
+    CloneMap.at(GEP)->setOperand(1, IV->getStartValue());
+
+  // Split vector preheader to form a new bypass block.
+  VPBasicBlock *NewPH = VectorPH->splitAt(VectorPH->end());
+  VPBasicBlock *ScalarPH = Plan.getScalarPreheader();
+
+  // Create bypass block branch.
+  VPRecipeBase *Uncountable = (*Cond)->getDefiningRecipe();
+  VPRecipeBase *PHUncountable = CloneMap.at(Uncountable);
+  VPBuilder PHBuilder(VectorPH, VectorPH->end());
+  VPValue *PHAnyOf = PHBuilder.createNaryOp(
+      VPInstruction::AnyOf, {PHUncountable->getVPSingleValue()});
+  PHBuilder.createNaryOp(VPInstruction::BranchOnCond, {PHAnyOf},
+                         PHUncountable->getDebugLoc());
+  VectorPH->clearSuccessors();
+  NewPH->clearPredecessors();
+  VPBlockUtils::connectBlocks(VectorPH, ScalarPH);
+  VPBlockUtils::connectBlocks(VectorPH, NewPH);
----------------
gbossu wrote:

I think this would deserve some more explaining. What does the control flow look like now?

I guess we start with
```
     VectorPH
        |
    VectorLoop
        |
 ...   ...
  \     |
   ScalarPH
```

And now have
```
VectorPH
       |    \
       |    NewPH
       |      |
       |   VectorLoop
       |      |
...    |     ...
 \     |     /
   ScalarPH
```
?

AFAIU, the preheader of the vector loop is now `NewPH`. Why is it important to override this and set it to `NewPH`?

https://github.com/llvm/llvm-project/pull/148626