[VPlan] Move predication to VPlanTransform (NFC) (WIP).

This patch moves the logic to predicate and linearize a VPlan to a
dedicated VPlan transform.

The main logic to perform predication is ready to review, although
there are few things to note that should be improved, either directly in
the PR or in the future:
 * Edge and block masks are cached in VPRecipeBuilder, so they can be
   accessed during recipe construction. A better alternative may be to
   add mask operands to all VPInstructions that need them and use that
   during recipe construction
 * The mask caching in a map also means that this map needs updating
   each time a new recipe replaces a VPInstruction; this would also be
   handled by adding mask operands.

Currently this is still WIP due to early-exit loop handling not working
due to the exit conditions not being available in the initial VPlans.
This will be fixed with #128419
and follow-ups

All tests except early-exit loops are passing

Author: fhahn

llvm/lib/Transforms/Vectorize/CMakeLists.txt

@@ -23,6 +23,7 @@ add_llvm_component_library(LLVMVectorize
   VPlan.cpp
   VPlanAnalysis.cpp
   VPlanHCFGBuilder.cpp
+  VPlanPredicator.cpp
   VPlanRecipes.cpp
   VPlanSLP.cpp
   VPlanTransforms.cpp

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

@@ -73,11 +73,14 @@ class VPRecipeBuilder {
   /// if-conversion currently takes place during VPlan-construction, so these
   /// caches are only used at that stage.
   using EdgeMaskCacheTy =
-      DenseMap<std::pair<BasicBlock *, BasicBlock *>, VPValue *>;
-  using BlockMaskCacheTy = DenseMap<BasicBlock *, VPValue *>;
+      DenseMap<std::pair<VPBasicBlock *, VPBasicBlock *>, VPValue *>;
+  using BlockMaskCacheTy = DenseMap<VPBasicBlock *, VPValue *>;
   EdgeMaskCacheTy EdgeMaskCache;
+
+public:
   BlockMaskCacheTy BlockMaskCache;
 
+private:
   // VPlan construction support: Hold a mapping from ingredients to
   // their recipe.
   DenseMap<Instruction *, VPRecipeBase *> Ingredient2Recipe;
@@ -114,11 +117,6 @@ class VPRecipeBuilder {
   tryToOptimizeInductionTruncate(TruncInst *I, ArrayRef<VPValue *> Operands,
                                  VFRange &Range);
 
-  /// Handle non-loop phi nodes. Return a new VPBlendRecipe otherwise. Currently
-  /// all such phi nodes are turned into a sequence of select instructions as
-  /// the vectorizer currently performs full if-conversion.
-  VPBlendRecipe *tryToBlend(PHINode *Phi, ArrayRef<VPValue *> Operands);
-
   /// Handle call instructions. If \p CI can be widened for \p Range.Start,
   /// return a new VPWidenCallRecipe or VPWidenIntrinsicRecipe. Range.End may be
   /// decreased to ensure same decision from \p Range.Start to \p Range.End.
@@ -187,27 +185,20 @@ class VPRecipeBuilder {
     Ingredient2Recipe[I] = R;
   }
 
-  /// Create the mask for the vector loop header block.
-  void createHeaderMask();
-
-  /// A helper function that computes the predicate of the block BB, assuming
-  /// that the header block of the loop is set to True or the loop mask when
-  /// tail folding.
-  void createBlockInMask(BasicBlock *BB);
-
+  void setBlockInMask(VPBasicBlock *BB, VPValue *Mask) {
+    assert(!BlockMaskCache.contains(BB) && "Mask already set");
+    BlockMaskCache[BB] = Mask;
+  }
   /// Returns the *entry* mask for the block \p BB.
-  VPValue *getBlockInMask(BasicBlock *BB) const;
-
-  /// Create an edge mask for every destination of cases and/or default.
-  void createSwitchEdgeMasks(SwitchInst *SI);
-
-  /// A helper function that computes the predicate of the edge between SRC
-  /// and DST.
-  VPValue *createEdgeMask(BasicBlock *Src, BasicBlock *Dst);
-
-  /// A helper that returns the previously computed predicate of the edge
-  /// between SRC and DST.
-  VPValue *getEdgeMask(BasicBlock *Src, BasicBlock *Dst) const;
+  VPValue *getBlockInMask(VPBasicBlock *BB) const {
+    return BlockMaskCache.lookup(BB);
+  }
+  void setEdgeMask(VPBasicBlock *Src, VPBasicBlock *Dst, VPValue *Mask) {
+    EdgeMaskCache[{Src, Dst}] = Mask;
+  }
+  VPValue *getEdgeMask(VPBasicBlock *Src, VPBasicBlock *Dst) const {
+    return EdgeMaskCache.lookup({Src, Dst});
+  }
 
   /// Return the recipe created for given ingredient.
   VPRecipeBase *getRecipe(Instruction *I) {

llvm/lib/Transforms/Vectorize/VPRecipeBuilder.h

@@ -75,7 +75,7 @@ class PlainCFGBuilder {
       : TheLoop(Lp), LI(LI), Plan(P) {}
 
   /// Build plain CFG for TheLoop  and connects it to Plan's entry.
-  void buildPlainCFG(DenseMap<VPBlockBase *, BasicBlock *> &VPB2IRBB);
+  void buildPlainCFG();
 };
 } // anonymous namespace
 
@@ -340,19 +340,24 @@ void PlainCFGBuilder::createVPInstructionsForVPBB(VPBasicBlock *VPBB,
       for (Value *Op : Inst->operands())
         VPOperands.push_back(getOrCreateVPOperand(Op));
 
-      // Build VPInstruction for any arbitrary Instruction without specific
-      // representation in VPlan.
-      NewR = cast<VPInstruction>(
-          VPIRBuilder.createNaryOp(Inst->getOpcode(), VPOperands, Inst));
+      if (auto *ICmp = dyn_cast<ICmpInst>(Inst)) {
+        NewR = cast<VPInstruction>(VPIRBuilder.createICmp(
+            ICmp->getPredicate(), VPOperands[0], VPOperands[1]));
+        NewR->setUnderlyingValue(ICmp);
+      } else {
+        // Build VPInstruction for any arbitrary Instruction without specific
+        // representation in VPlan.
+        NewR = cast<VPInstruction>(
+            VPIRBuilder.createNaryOp(Inst->getOpcode(), VPOperands, Inst));
+      }
     }
 
     IRDef2VPValue[Inst] = NewR;
   }
 }
 
 // Main interface to build the plain CFG.
-void PlainCFGBuilder::buildPlainCFG(
-    DenseMap<VPBlockBase *, BasicBlock *> &VPB2IRBB) {
+void PlainCFGBuilder::buildPlainCFG() {
   // 0. Reuse the top-level region, vector-preheader and exit VPBBs from the
   // skeleton. These were created directly rather than via getOrCreateVPBB(),
   // revisit them now to update BB2VPBB. Note that header/entry and
@@ -481,14 +486,11 @@ void PlainCFGBuilder::buildPlainCFG(
   // have a VPlan counterpart. Fix VPlan header phi by adding their
   // corresponding VPlan operands.
   fixHeaderPhis();
-
-  for (const auto &[IRBB, VPB] : BB2VPBB)
-    VPB2IRBB[VPB] = IRBB;
 }
 
 void VPlanHCFGBuilder::buildPlainCFG() {
   PlainCFGBuilder PCFGBuilder(TheLoop, LI, Plan);
-  PCFGBuilder.buildPlainCFG(VPB2IRBB);
+  PCFGBuilder.buildPlainCFG();
 }
 
 // Public interface to build a H-CFG.

llvm/lib/Transforms/Vectorize/VPlanHCFGBuilder.cpp

@@ -53,10 +53,6 @@ class VPlanHCFGBuilder {
   // are introduced.
   VPDominatorTree VPDomTree;
 
-  /// Map of create VP blocks to their input IR basic blocks, if they have been
-  /// created for a input IR basic block.
-  DenseMap<VPBlockBase *, BasicBlock *> VPB2IRBB;
-
   /// Build plain CFG for TheLoop and connects it to Plan's entry.
   void buildPlainCFG();
 
@@ -66,14 +62,6 @@ class VPlanHCFGBuilder {
 
   /// Build H-CFG for TheLoop and update Plan accordingly.
   void buildHierarchicalCFG();
-
-  /// Return the input IR BasicBlock corresponding to \p VPB. Returns nullptr if
-  /// there is no such corresponding block.
-  /// FIXME: This is a temporary workaround to drive the createBlockInMask.
-  /// Remove once mask creation is done on VPlan.
-  BasicBlock *getIRBBForVPB(const VPBlockBase *VPB) const {
-    return VPB2IRBB.lookup(VPB);
-  }
 };
 } // namespace llvm