Revert "[VPlan] Move predication to VPlanTransform (NFC). (#128420)"

This reverts commit b263c08.

Looks like this triggers a crash in one of the Fortran tests. Reverting
while I investigate
    https://lab.llvm.org/buildbot/#/builders/41/builds/6825

Author: fhahn

llvm/lib/Transforms/Vectorize/CMakeLists.txt

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

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

@@ -68,10 +68,15 @@ class VPRecipeBuilder {
 
   VPBuilder &Builder;
 
-  /// The mask of each VPBB, generated earlier and used for predicating recipes
-  /// in VPBB.
-  /// TODO: remove by applying predication when generating the masks.
-  DenseMap<VPBasicBlock *, VPValue *> &BlockMaskCache;
+  /// When we if-convert we need to create edge masks. We have to cache values
+  /// so that we don't end up with exponential recursion/IR. Note that
+  /// 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 *>;
+  EdgeMaskCacheTy EdgeMaskCache;
+  BlockMaskCacheTy BlockMaskCache;
 
   // VPlan construction support: Hold a mapping from ingredients to
   // their recipe.
@@ -85,6 +90,10 @@ class VPRecipeBuilder {
   /// A mapping of partial reduction exit instructions to their scaling factor.
   DenseMap<const Instruction *, unsigned> ScaledReductionMap;
 
+  /// A mapping from VP blocks to IR blocks, used temporarily while migrating
+  /// away from IR references.
+  const DenseMap<const VPBlockBase *, BasicBlock *> &VPB2IRBB;
+
   /// Loop versioning instance for getting noalias metadata guaranteed by
   /// runtime checks.
   LoopVersioning *LVer;
@@ -113,6 +122,11 @@ class VPRecipeBuilder {
   tryToOptimizeInductionTruncate(TruncInst *I, ArrayRef<VPValue *> Operands,
                                  VFRange &Range);
 
+  /// Handle non-loop phi nodes, returning a new VPBlendRecipe. Currently
+  /// all such phi nodes are turned into a sequence of select instructions as
+  /// the vectorizer currently performs full if-conversion.
+  VPBlendRecipe *tryToBlend(VPWidenPHIRecipe *PhiR);
+
   /// 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.
@@ -150,11 +164,10 @@ class VPRecipeBuilder {
                   LoopVectorizationLegality *Legal,
                   LoopVectorizationCostModel &CM,
                   PredicatedScalarEvolution &PSE, VPBuilder &Builder,
-                  DenseMap<VPBasicBlock *, VPValue *> &BlockMaskCache,
+                  const DenseMap<const VPBlockBase *, BasicBlock *> &VPB2IRBB,
                   LoopVersioning *LVer)
       : Plan(Plan), OrigLoop(OrigLoop), TLI(TLI), TTI(TTI), Legal(Legal),
-        CM(CM), PSE(PSE), Builder(Builder), BlockMaskCache(BlockMaskCache),
-        LVer(LVer) {}
+        CM(CM), PSE(PSE), Builder(Builder), VPB2IRBB(VPB2IRBB), LVer(LVer) {}
 
   std::optional<unsigned> getScalingForReduction(const Instruction *ExitInst) {
     auto It = ScaledReductionMap.find(ExitInst);
@@ -183,11 +196,38 @@ class VPRecipeBuilder {
     Ingredient2Recipe[I] = R;
   }
 
-  /// Returns the *entry* mask for block \p VPBB or null if the mask is
-  /// all-true.
-  VPValue *getBlockInMask(VPBasicBlock *VPBB) const {
-    return BlockMaskCache.lookup(VPBB);
+  /// 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(const VPBasicBlock *VPBB) {
+    return createBlockInMask(VPB2IRBB.lookup(VPBB));
   }
+  void createBlockInMask(BasicBlock *BB);
+
+  /// Returns the *entry* mask for the block \p VPBB.
+  VPValue *getBlockInMask(const VPBasicBlock *VPBB) const {
+    return getBlockInMask(VPB2IRBB.lookup(VPBB));
+  }
+
+  /// 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(const VPBasicBlock *Src, const VPBasicBlock *Dst) const {
+    return getEdgeMask(VPB2IRBB.lookup(Src), VPB2IRBB.lookup(Dst));
+  }
+  VPValue *getEdgeMask(BasicBlock *Src, BasicBlock *Dst) const;
 
   /// Return the recipe created for given ingredient.
   VPRecipeBase *getRecipe(Instruction *I) {
@@ -212,15 +252,6 @@ class VPRecipeBuilder {
     }
     return Plan.getOrAddLiveIn(V);
   }
-
-  void updateBlockMaskCache(DenseMap<VPValue *, VPValue *> &Old2New) {
-    for (auto &[_, V] : BlockMaskCache) {
-      if (auto *New = Old2New.lookup(V)) {
-        V->replaceAllUsesWith(New);
-        V = New;
-      }
-    }
-  }
 };
 } // end namespace llvm
 

llvm/lib/Transforms/Vectorize/VPRecipeBuilder.h

@@ -65,8 +65,9 @@ class PlainCFGBuilder {
   PlainCFGBuilder(Loop *Lp, LoopInfo *LI)
       : TheLoop(Lp), LI(LI), Plan(std::make_unique<VPlan>(Lp)) {}
 
-  /// Build plain CFG for TheLoop and connect it to Plan's entry.
-  std::unique_ptr<VPlan> buildPlainCFG();
+  /// Build plain CFG for TheLoop  and connects it to Plan's entry.
+  std::unique_ptr<VPlan>
+  buildPlainCFG(DenseMap<const VPBlockBase *, BasicBlock *> &VPB2IRBB);
 };
 } // anonymous namespace
 
@@ -241,7 +242,8 @@ void PlainCFGBuilder::createVPInstructionsForVPBB(VPBasicBlock *VPBB,
 }
 
 // Main interface to build the plain CFG.
-std::unique_ptr<VPlan> PlainCFGBuilder::buildPlainCFG() {
+std::unique_ptr<VPlan> PlainCFGBuilder::buildPlainCFG(
+    DenseMap<const VPBlockBase *, BasicBlock *> &VPB2IRBB) {
   VPIRBasicBlock *Entry = cast<VPIRBasicBlock>(Plan->getEntry());
   BB2VPBB[Entry->getIRBasicBlock()] = Entry;
   for (VPIRBasicBlock *ExitVPBB : Plan->getExitBlocks())
@@ -332,14 +334,18 @@ std::unique_ptr<VPlan> PlainCFGBuilder::buildPlainCFG() {
     }
   }
 
+  for (const auto &[IRBB, VPB] : BB2VPBB)
+    VPB2IRBB[VPB] = IRBB;
+
   LLVM_DEBUG(Plan->setName("Plain CFG\n"); dbgs() << *Plan);
   return std::move(Plan);
 }
 
-std::unique_ptr<VPlan> VPlanTransforms::buildPlainCFG(Loop *TheLoop,
-                                                      LoopInfo &LI) {
+std::unique_ptr<VPlan> VPlanTransforms::buildPlainCFG(
+    Loop *TheLoop, LoopInfo &LI,
+    DenseMap<const VPBlockBase *, BasicBlock *> &VPB2IRBB) {
   PlainCFGBuilder Builder(TheLoop, &LI);
-  return Builder.buildPlainCFG();
+  return Builder.buildPlainCFG(VPB2IRBB);
 }
 
 /// Checks if \p HeaderVPB is a loop header block in the plain CFG; that is, it