[VPlan] Add VPlan::hasTailFolded.

[fhahn]

Summary

This patch adds a helper to check if the tail loop is folded in a
VPlan.

To do so it checks if the VPlan contains a widen canonical IV, a
active-lane-mask-phi, an active-lane-mask or a ICmpULE compare of a
widened canonical induction with the backedge taken count. This should
effectively check if a mask for the header is present, but may be a bit
cumbersome.

Alternatively we could also add a flag to the VPlan.

Depends on D157037.

[Ayal]

Alternatively we could also add a flag to the VPlan.

Raising thoughts of other potential alternatives.

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

3917

Could recipes be set initially such that they generate code as needed, rather than (figuring out if tail was folded in order to) fix it here? This seems to require (a) setting FMF of Sel's recipe, (b) optionally rewiring a header phi recipe to Sel's recipe, (c) RAUW of LoopExitInstDef with Sel's recipe, rather than State.reset().

9259

Here, during VPlan construction after header mask was proactively created (depending on D157037), it would be easy/easier to ask RecipeBuilder if the mask cached for the header block is null? Namely, check if Cond below is null.

9271

Set the FMF of the select recipe upon creation here, rather than fixing its generated IR in fixReduction()?

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

void InnerLoopVectorizer::fixReduction(VPReductionPHIRecipe *PhiR,
 
   VPBasicBlock *LatchVPBB =
       PhiR->getParent()->getEnclosingLoopRegion()->getExitingBasicBlock();
   BasicBlock *VectorLoopLatch = State.CFG.VPBB2IRBB[LatchVPBB];
   // If tail is folded by masking, the vector value to leave the loop should be
   // a Select choosing between the vectorized LoopExitInst and vectorized Phi,
   // instead of the former. For an inloop reduction the reduction will already
   // be predicated, and does not need to be handled here.
-  if (Cost->foldTailByMasking() && !PhiR->isInLoop()) {
+  if (PhiR->getParent()->getPlan()->hasTailFolded() && !PhiR->isInLoop()) {
     for (unsigned Part = 0; Part < UF; ++Part) {
       Value *VecLoopExitInst = State.get(LoopExitInstDef, Part);
       SelectInst *Sel = nullptr;
       for (User *U : VecLoopExitInst->users()) {
         if (isa<SelectInst>(U)) {
           assert(!Sel && "Reduction exit feeding two selects");
           Sel = cast<SelectInst>(U);
         } else
for (VPRecipeBase *CurrentLink : Worklist.getArrayRef().drop_front()) {
       CurrentLink->getVPSingleValue()->replaceAllUsesWith(RedRecipe);
       PreviousLink = RedRecipe;
     }
   }
 
   // If tail is folded by masking, introduce selects between the phi
   // and the live-out instruction of each reduction, at the beginning of the
   // dedicated latch block.
-  if (CM.foldTailByMasking()) {
+  if (Plan->hasTailFolded()) {
     Builder.setInsertPoint(LatchVPBB, LatchVPBB->begin());
     for (VPRecipeBase &R :
          Plan->getVectorLoopRegion()->getEntryBasicBlock()->phis()) {
       VPReductionPHIRecipe *PhiR = dyn_cast<VPReductionPHIRecipe>(&R);
       if (!PhiR || PhiR->isInLoop())
         continue;
       VPValue *Cond =
           RecipeBuilder.createBlockInMask(OrigLoop->getHeader(), *Plan);
       VPValue *Red = PhiR->getBackedgeValue();
       assert(Red->getDefiningRecipe()->getParent() != LatchVPBB &&
              "reduction recipe must be defined before latch");
       Builder.createNaryOp(Instruction::Select, {Cond, Red, PhiR});
     }
   }
 
   VPlanTransforms::clearReductionWrapFlags(*Plan);
 }
 
 #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
 void VPInterleaveRecipe::print(raw_ostream &O, const Twine &Indent,

llvm/lib/Transforms/Vectorize/VPlan.h

void addSCEVExpansion(const SCEV *S, VPValue *V) {
     assert(!SCEVToExpansion.contains(S) && "SCEV already expanded");
     SCEVToExpansion[S] = V;
   }
 
   /// \return The block corresponding to the original preheader.
   VPBasicBlock *getPreheader() { return Preheader; }
   const VPBasicBlock *getPreheader() const { return Preheader; }
 
+  /// Returns true if the tail loop is folded in the VPlan.
+  bool hasTailFolded();
+
 private:
   /// Add to the given dominator tree the header block and every new basic block
   /// that was created between it and the latch block, inclusive.
   static void updateDominatorTree(DominatorTree *DT, BasicBlock *LoopLatchBB,
                                   BasicBlock *LoopPreHeaderBB,
                                   BasicBlock *LoopExitBB);
 };
 

llvm/lib/Transforms/Vectorize/VPlan.cpp

 void VPlan::dump() const { print(dbgs()); }
 #endif
 
 void VPlan::addLiveOut(PHINode *PN, VPValue *V) {
   assert(LiveOuts.count(PN) == 0 && "an exit value for PN already exists");
   LiveOuts.insert({PN, new VPLiveOut(PN, V)});
 }
 
+bool VPlan::hasTailFolded() {
+  VPBasicBlock *HeaderVPBB = getVectorLoopRegion()->getEntryBasicBlock();
+  return any_of(*HeaderVPBB, [this](VPRecipeBase &R) {
+    // A widened canonical IV or active-lane-mask PHI imply the tail being
+    // folded.
+    if (isa<VPWidenCanonicalIVRecipe>(&R) || isa<VPActiveLaneMaskPHIRecipe>(&R))
+      return true;
+
+    auto *VPI = dyn_cast<VPInstruction>(&R);
+    if (!VPI)
+      return false;
+    if (VPI->getOpcode() == VPInstruction::ActiveLaneMask)
+      return true;
+    if (VPI->getOpcode() != VPInstruction::ICmpULE)
+      return false;
+    auto IsWidenCanonicalIV = [](const VPRecipeBase &R) {
+      if (auto *Ind = dyn_cast<VPWidenIntOrFpInductionRecipe>(&R))
+        return Ind->isCanonical();
+      return isa<VPWidenCanonicalIVRecipe>(&R);
+    };
+    if (IsWidenCanonicalIV(*VPI->getOperand(0)->getDefiningRecipe()) &&
+        VPI->getOperand(1) == getOrCreateBackedgeTakenCount())
+      return true;
+  });
+}
+
 void VPlan::updateDominatorTree(DominatorTree *DT, BasicBlock *LoopHeaderBB,
                                 BasicBlock *LoopLatchBB,
                                 BasicBlock *LoopExitBB) {
   // The vector body may be more than a single basic-block by this point.
   // Update the dominator tree information inside the vector body by propagating
   // it from header to latch, expecting only triangular control-flow, if any.
   BasicBlock *PostDomSucc = nullptr;
   for (auto *BB = LoopHeaderBB; BB != LoopLatchBB; BB = PostDomSucc) {