[VPlan] Model entry-mask of VPBB in VPlan (NFC)

Model the entry-mask of a VPBB in VPlan, replacing the Old2New and
BlockMasksCache hacks with a simple RAUW. It is observed that the
entry-mask of a VPBB is unused after adjustRecipesForReductions: hence,
erase all the masks past this transform, avoiding a wide-scale refactor
to take non-recipe VPUsers into account.

Author: artagnon

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

@@ -7578,7 +7578,7 @@ VPWidenMemoryRecipe *VPRecipeBuilder::tryToWidenMemory(VPInstruction *VPI,
 
   VPValue *Mask = nullptr;
   if (Legal->isMaskRequired(I))
-    Mask = getBlockInMask(Builder.getInsertBlock());
+    Mask = Builder.getInsertBlock()->getEntryMask();
 
   // Determine if the pointer operand of the access is either consecutive or
   // reverse consecutive.
@@ -7791,7 +7791,7 @@ VPSingleDefRecipe *VPRecipeBuilder::tryToWidenCall(VPInstruction *VPI,
       //      all-true mask.
       VPValue *Mask = nullptr;
       if (Legal->isMaskRequired(CI))
-        Mask = getBlockInMask(Builder.getInsertBlock());
+        Mask = Builder.getInsertBlock()->getEntryMask();
       else
         Mask = Plan.getOrAddLiveIn(
             ConstantInt::getTrue(IntegerType::getInt1Ty(Plan.getContext())));
@@ -7834,7 +7834,7 @@ VPWidenRecipe *VPRecipeBuilder::tryToWiden(VPInstruction *VPI) {
     // div/rem operation itself.  Otherwise fall through to general handling below.
     if (CM.isPredicatedInst(I)) {
       SmallVector<VPValue *> Ops(VPI->operands());
-      VPValue *Mask = getBlockInMask(Builder.getInsertBlock());
+      VPValue *Mask = Builder.getInsertBlock()->getEntryMask();
       VPValue *One = Plan.getConstantInt(I->getType(), 1u);
       auto *SafeRHS =
           Builder.createSelect(Mask, Ops[1], One, VPI->getDebugLoc());
@@ -7914,7 +7914,7 @@ VPHistogramRecipe *VPRecipeBuilder::tryToWidenHistogram(const HistogramInfo *HI,
   // In case of predicated execution (due to tail-folding, or conditional
   // execution, or both), pass the relevant mask.
   if (Legal->isMaskRequired(HI->Store))
-    HGramOps.push_back(getBlockInMask(Builder.getInsertBlock()));
+    HGramOps.push_back(Builder.getInsertBlock()->getEntryMask());
 
   return new VPHistogramRecipe(Opcode, HGramOps, VPI->getDebugLoc());
 }
@@ -7968,7 +7968,7 @@ VPReplicateRecipe *VPRecipeBuilder::handleReplication(VPInstruction *VPI,
     // added initially. Masked replicate recipes will later be placed under an
     // if-then construct to prevent side-effects. Generate recipes to compute
     // the block mask for this region.
-    BlockInMask = getBlockInMask(Builder.getInsertBlock());
+    BlockInMask = Builder.getInsertBlock()->getEntryMask();
   }
 
   // Note that there is some custom logic to mark some intrinsics as uniform
@@ -8302,7 +8302,7 @@ VPRecipeBuilder::tryToCreatePartialReduction(VPInstruction *Reduction,
 
   VPValue *Cond = nullptr;
   if (CM.blockNeedsPredicationForAnyReason(ReductionI->getParent()))
-    Cond = getBlockInMask(Builder.getInsertBlock());
+    Cond = Builder.getInsertBlock()->getEntryMask();
   return new VPPartialReductionRecipe(ReductionOpcode, Accumulator, BinOp, Cond,
                                       ScaleFactor, ReductionI);
 }
@@ -8426,15 +8426,14 @@ VPlanPtr LoopVectorizationPlanner::tryToBuildVPlanWithVPRecipes(
   // ---------------------------------------------------------------------------
   // Predicate and linearize the top-level loop region.
   // ---------------------------------------------------------------------------
-  auto BlockMaskCache = VPlanTransforms::introduceMasksAndLinearize(
-      *Plan, CM.foldTailByMasking());
+  VPlanTransforms::introduceMasksAndLinearize(*Plan, CM.foldTailByMasking());
 
   // ---------------------------------------------------------------------------
   // Construct wide recipes and apply predication for original scalar
   // VPInstructions in the loop.
   // ---------------------------------------------------------------------------
   VPRecipeBuilder RecipeBuilder(*Plan, OrigLoop, TLI, &TTI, Legal, CM, PSE,
-                                Builder, BlockMaskCache);
+                                Builder);
   // TODO: Handle partial reductions with EVL tail folding.
   if (!CM.foldTailWithEVL())
     RecipeBuilder.collectScaledReductions(Range);
@@ -8447,9 +8446,6 @@ VPlanPtr LoopVectorizationPlanner::tryToBuildVPlanWithVPRecipes(
 
   auto *MiddleVPBB = Plan->getMiddleBlock();
   VPBasicBlock::iterator MBIP = MiddleVPBB->getFirstNonPhi();
-  // Mapping from VPValues in the initial plan to their widened VPValues. Needed
-  // temporarily to update created block masks.
-  DenseMap<VPValue *, VPValue *> Old2New;
   for (VPBasicBlock *VPBB : VPBlockUtils::blocksOnly<VPBasicBlock>(RPOT)) {
     // Convert input VPInstructions to widened recipes.
     for (VPRecipeBase &R : make_early_inc_range(*VPBB)) {
@@ -8505,7 +8501,7 @@ VPlanPtr LoopVectorizationPlanner::tryToBuildVPlanWithVPRecipes(
       }
       if (Recipe->getNumDefinedValues() == 1) {
         SingleDef->replaceAllUsesWith(Recipe->getVPSingleValue());
-        Old2New[SingleDef] = Recipe->getVPSingleValue();
+        SingleDef->eraseFromParent();
       } else {
         assert(Recipe->getNumDefinedValues() == 0 &&
                "Unexpected multidef recipe");
@@ -8514,14 +8510,6 @@ VPlanPtr LoopVectorizationPlanner::tryToBuildVPlanWithVPRecipes(
     }
   }
 
-  // replaceAllUsesWith above may invalidate the block masks. Update them here.
-  // TODO: Include the masks as operands in the predicated VPlan directly
-  // to remove the need to keep a map of masks beyond the predication
-  // transform.
-  RecipeBuilder.updateBlockMaskCache(Old2New);
-  for (VPValue *Old : Old2New.keys())
-    Old->getDefiningRecipe()->eraseFromParent();
-
   assert(isa<VPRegionBlock>(LoopRegion) &&
          !LoopRegion->getEntryBasicBlock()->empty() &&
          "entry block must be set to a VPRegionBlock having a non-empty entry "
@@ -8541,6 +8529,11 @@ VPlanPtr LoopVectorizationPlanner::tryToBuildVPlanWithVPRecipes(
   // Adjust the recipes for any inloop reductions.
   adjustRecipesForReductions(Plan, RecipeBuilder, Range.Start);
 
+  // Erase the block entry masks, since they're not used any longer, so that
+  // future transforms only deal with recipe VPUsers.
+  for (VPBasicBlock *VPBB : VPBlockUtils::blocksOnly<VPBasicBlock>(RPOT))
+    VPBB->eraseEntryMask();
+
   // Apply mandatory transformation to handle FP maxnum/minnum reduction with
   // NaNs if possible, bail out otherwise.
   if (!VPlanTransforms::runPass(VPlanTransforms::handleMaxMinNumReductions,
@@ -8792,7 +8785,7 @@ void LoopVectorizationPlanner::adjustRecipesForReductions(
 
       VPValue *CondOp = nullptr;
       if (CM.blockNeedsPredicationForAnyReason(CurrentLinkI->getParent()))
-        CondOp = RecipeBuilder.getBlockInMask(CurrentLink->getParent());
+        CondOp = CurrentLink->getParent()->getEntryMask();
 
       auto *RedRecipe = new VPReductionRecipe(
           Kind, FMFs, CurrentLinkI, PreviousLink, VecOp, CondOp,
@@ -8833,7 +8826,7 @@ void LoopVectorizationPlanner::adjustRecipesForReductions(
     // different numbers of lanes. Partial reductions mask the input instead.
     if (!PhiR->isInLoop() && CM.foldTailByMasking() &&
         !isa<VPPartialReductionRecipe>(OrigExitingVPV)) {
-      VPValue *Cond = RecipeBuilder.getBlockInMask(PhiR->getParent());
+      VPValue *Cond = PhiR->getParent()->getEntryMask();
       std::optional<FastMathFlags> FMFs =
           PhiTy->isFloatingPointTy()
               ? std::make_optional(RdxDesc.getFastMathFlags())

llvm/lib/Transforms/Vectorize/VPRecipeBuilder.h

@@ -67,11 +67,6 @@ 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;
-
   // VPlan construction support: Hold a mapping from ingredients to
   // their recipe.
   DenseMap<Instruction *, VPRecipeBase *> Ingredient2Recipe;
@@ -138,10 +133,9 @@ class VPRecipeBuilder {
                   const TargetTransformInfo *TTI,
                   LoopVectorizationLegality *Legal,
                   LoopVectorizationCostModel &CM,
-                  PredicatedScalarEvolution &PSE, VPBuilder &Builder,
-                  DenseMap<VPBasicBlock *, VPValue *> &BlockMaskCache)
+                  PredicatedScalarEvolution &PSE, VPBuilder &Builder)
       : Plan(Plan), OrigLoop(OrigLoop), TLI(TLI), TTI(TTI), Legal(Legal),
-        CM(CM), PSE(PSE), Builder(Builder), BlockMaskCache(BlockMaskCache) {}
+        CM(CM), PSE(PSE), Builder(Builder) {}
 
   std::optional<unsigned> getScalingForReduction(const Instruction *ExitInst) {
     auto It = ScaledReductionMap.find(ExitInst);
@@ -169,12 +163,6 @@ 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);
-  }
-
   /// Return the recipe created for given ingredient.
   VPRecipeBase *getRecipe(Instruction *I) {
     assert(Ingredient2Recipe.count(I) &&
@@ -196,15 +184,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/VPlan.h

@@ -3982,13 +3982,15 @@ struct CastInfo<VPIRMetadata, const VPRecipeBase *>
 
 /// VPBasicBlock serves as the leaf of the Hierarchical Control-Flow Graph. It
 /// holds a sequence of zero or more VPRecipe's each representing a sequence of
-/// output IR instructions. All PHI-like recipes must come before any non-PHI recipes.
-class LLVM_ABI_FOR_TEST VPBasicBlock : public VPBlockBase {
+/// output IR instructions. All PHI-like recipes must come before any non-PHI
+/// recipes. It also has an operand corresponding to a mask on which to enter
+/// the block, which is used early in the VPlan construction.
+class LLVM_ABI_FOR_TEST VPBasicBlock : public VPBlockBase, protected VPUser {
   friend class VPlan;
 
   /// Use VPlan::createVPBasicBlock to create VPBasicBlocks.
   VPBasicBlock(const Twine &Name = "", VPRecipeBase *Recipe = nullptr)
-      : VPBlockBase(VPBasicBlockSC, Name.str()) {
+      : VPBlockBase(VPBasicBlockSC, Name.str()), VPUser(VPUBlockSC) {
     if (Recipe)
       appendRecipe(Recipe);
   }
@@ -4001,7 +4003,7 @@ class LLVM_ABI_FOR_TEST VPBasicBlock : public VPBlockBase {
   RecipeListTy Recipes;
 
   VPBasicBlock(const unsigned char BlockSC, const Twine &Name = "")
-      : VPBlockBase(BlockSC, Name.str()) {}
+      : VPBlockBase(BlockSC, Name.str()), VPUser(VPUBlockSC) {}
 
 public:
   ~VPBasicBlock() override {
@@ -4112,6 +4114,21 @@ class LLVM_ABI_FOR_TEST VPBasicBlock : public VPBlockBase {
   /// second predecessor is the exiting block of the region.
   const VPBasicBlock *getCFGPredecessor(unsigned Idx) const;
 
+  /// Get the entry mask of this block. nullptr is used to communicate an
+  /// all-ones mask.
+  VPValue *getEntryMask() const {
+    return getNumOperands() ? getOperand(0) : nullptr;
+  }
+
+  /// Set the entry mask of this block: used by VPlanPredicator, when
+  /// predicating blocks.
+  void setEntryMask(VPValue *M) {
+    getNumOperands() ? setOperand(0, M) : addOperand(M);
+  }
+
+  /// Erase the entry mask of this block.
+  void eraseEntryMask() { eraseOperands(); }
+
 protected:
   /// Execute the recipes in the IR basic block \p BB.
   void executeRecipes(VPTransformState *State, BasicBlock *BB);

llvm/lib/Transforms/Vectorize/VPlanPredicator.cpp

@@ -32,32 +32,15 @@ class VPPredicator {
   using EdgeMaskCacheTy =
       DenseMap<std::pair<const VPBasicBlock *, const VPBasicBlock *>,
                VPValue *>;
-  using BlockMaskCacheTy = DenseMap<VPBasicBlock *, VPValue *>;
   EdgeMaskCacheTy EdgeMaskCache;
 
-  BlockMaskCacheTy BlockMaskCache;
-
   /// Create an edge mask for every destination of cases and/or default.
   void createSwitchEdgeMasks(VPInstruction *SI);
 
   /// Computes and return the predicate of the edge between \p Src and \p Dst,
   /// possibly inserting new recipes at \p Dst (using Builder's insertion point)
   VPValue *createEdgeMask(VPBasicBlock *Src, VPBasicBlock *Dst);
 
-  /// Returns the *entry* mask for \p VPBB.
-  VPValue *getBlockInMask(VPBasicBlock *VPBB) const {
-    return BlockMaskCache.lookup(VPBB);
-  }
-
-  /// Record \p Mask as the *entry* mask of \p VPBB, which is expected to not
-  /// already have a mask.
-  void setBlockInMask(VPBasicBlock *VPBB, VPValue *Mask) {
-    // TODO: Include the masks as operands in the predicated VPlan directly to
-    // avoid keeping the map of masks beyond the predication transform.
-    assert(!getBlockInMask(VPBB) && "Mask already set");
-    BlockMaskCache[VPBB] = Mask;
-  }
-
   /// Record \p Mask as the mask of the edge from \p Src to \p Dst. The edge is
   /// expected to not have a mask already.
   VPValue *setEdgeMask(const VPBasicBlock *Src, const VPBasicBlock *Dst,
@@ -82,8 +65,6 @@ class VPPredicator {
 
   /// Convert phi recipes in \p VPBB to VPBlendRecipes.
   void convertPhisToBlends(VPBasicBlock *VPBB);
-
-  const BlockMaskCacheTy getBlockMaskCache() const { return BlockMaskCache; }
 };
 } // namespace
 
@@ -95,7 +76,7 @@ VPValue *VPPredicator::createEdgeMask(VPBasicBlock *Src, VPBasicBlock *Dst) {
   if (EdgeMask)
     return EdgeMask;
 
-  VPValue *SrcMask = getBlockInMask(Src);
+  VPValue *SrcMask = Src->getEntryMask();
 
   // If there's a single successor, there's no terminator recipe.
   if (Src->getNumSuccessors() == 1)
@@ -140,7 +121,6 @@ VPValue *VPPredicator::createBlockInMask(VPBasicBlock *VPBB) {
     VPValue *EdgeMask = createEdgeMask(cast<VPBasicBlock>(Predecessor), VPBB);
     if (!EdgeMask) { // Mask of predecessor is all-one so mask of block is
                      // too.
-      setBlockInMask(VPBB, EdgeMask);
       return EdgeMask;
     }
 
@@ -152,15 +132,13 @@ VPValue *VPPredicator::createBlockInMask(VPBasicBlock *VPBB) {
     BlockMask = Builder.createOr(BlockMask, EdgeMask, {});
   }
 
-  setBlockInMask(VPBB, BlockMask);
+  VPBB->setEntryMask(BlockMask);
   return BlockMask;
 }
 
 void VPPredicator::createHeaderMask(VPBasicBlock *HeaderVPBB, bool FoldTail) {
-  if (!FoldTail) {
-    setBlockInMask(HeaderVPBB, nullptr);
+  if (!FoldTail)
     return;
-  }
 
   // Introduce the early-exit compare IV <= BTC to form header block mask.
   // This is used instead of IV < TC because TC may wrap, unlike BTC. Start by
@@ -175,7 +153,7 @@ void VPPredicator::createHeaderMask(VPBasicBlock *HeaderVPBB, bool FoldTail) {
 
   VPValue *BTC = Plan.getOrCreateBackedgeTakenCount();
   VPValue *BlockMask = Builder.createICmp(CmpInst::ICMP_ULE, IV, BTC);
-  setBlockInMask(HeaderVPBB, BlockMask);
+  HeaderVPBB->setEntryMask(BlockMask);
 }
 
 void VPPredicator::createSwitchEdgeMasks(VPInstruction *SI) {
@@ -201,7 +179,7 @@ void VPPredicator::createSwitchEdgeMasks(VPInstruction *SI) {
 
   // We need to handle 2 separate cases below for all entries in Dst2Compares,
   // which excludes destinations matching the default destination.
-  VPValue *SrcMask = getBlockInMask(Src);
+  VPValue *SrcMask = Src->getEntryMask();
   VPValue *DefaultMask = nullptr;
   for (const auto &[Dst, Conds] : Dst2Compares) {
     // 1. Dst is not the default destination. Dst is reached if any of the
@@ -261,8 +239,7 @@ void VPPredicator::convertPhisToBlends(VPBasicBlock *VPBB) {
   }
 }
 
-DenseMap<VPBasicBlock *, VPValue *>
-VPlanTransforms::introduceMasksAndLinearize(VPlan &Plan, bool FoldTail) {
+void VPlanTransforms::introduceMasksAndLinearize(VPlan &Plan, bool FoldTail) {
   VPRegionBlock *LoopRegion = Plan.getVectorLoopRegion();
   // Scan the body of the loop in a topological order to visit each basic block
   // after having visited its predecessor basic blocks.
@@ -301,5 +278,4 @@ VPlanTransforms::introduceMasksAndLinearize(VPlan &Plan, bool FoldTail) {
 
     PrevVPBB = VPBB;
   }
-  return Predicator.getBlockMaskCache();
 }

llvm/lib/Transforms/Vectorize/VPlanTransforms.h

@@ -360,11 +360,8 @@ struct VPlanTransforms {
   /// Predicate and linearize the control-flow in the only loop region of
   /// \p Plan. If \p FoldTail is true, create a mask guarding the loop
   /// header, otherwise use all-true for the header mask. Masks for blocks are
-  /// added to a block-to-mask map which is returned in order to be used later
-  /// for wide recipe construction. This argument is temporary and will be
-  /// removed in the future.
-  static DenseMap<VPBasicBlock *, VPValue *>
-  introduceMasksAndLinearize(VPlan &Plan, bool FoldTail);
+  /// added to blocks themselves.
+  static void introduceMasksAndLinearize(VPlan &Plan, bool FoldTail);
 
   /// Add branch weight metadata, if the \p Plan's middle block is terminated by
   /// a BranchOnCond recipe.