!fixup address comments, thanks

llvm · fhahn · May 21, 2025 · Feb 22, 2025 · May 16, 2025 · May 16, 2025
commit 58c8fc4992731e6f56a9af3aab9cef6f8d2f0a2c
diff --git a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
@@ -9254,6 +9254,9 @@ LoopVectorizationPlanner::tryToBuildVPlanWithVPRecipes(VFRange &Range,

  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)) {
@@ -9313,19 +9316,23 @@ LoopVectorizationPlanner::tryToBuildVPlanWithVPRecipes(VFRange &Range,
      }
      if (Recipe->getNumDefinedValues() == 1) {
        SingleDef->replaceAllUsesWith(Recipe->getVPSingleValue());
-        // replaceAllUsesWith may invalidate the block mask cache. Update it.
-        // 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(SingleDef,
-                                           Recipe->getVPSingleValue());
-      } else
+        Old2New[SingleDef] = Recipe->getVPSingleValue();
+      } else {
        assert(Recipe->getNumDefinedValues() == 0 &&
               "Unexpected multidef recipe");
-      R.eraseFromParent();
+        R.eraseFromParent();
+      }
    }
  }

+  // 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 (const auto &[Old, New] : Old2New)
-  for (const auto &[Old, New] : Old2New)
+  for (const auto &[Old, _] : Old2New)
-  for (const auto &[Old, New] : Old2New)
+  for (const auto &[Old, _] : Old2New)
+    Old->getDefiningRecipe()->eraseFromParent();
+
  assert(isa<VPRegionBlock>(Plan->getVectorLoopRegion()) &&
         !Plan->getVectorLoopRegion()->getEntryBasicBlock()->empty() &&
         "entry block must be set to a VPRegionBlock having a non-empty entry "

diff --git a/llvm/lib/Transforms/Vectorize/VPRecipeBuilder.h b/llvm/lib/Transforms/Vectorize/VPRecipeBuilder.h
@@ -209,10 +209,12 @@ class VPRecipeBuilder {
    return Plan.getOrAddLiveIn(V);
  }

-  void updateBlockMaskCache(VPValue *Old, VPValue *New) {
+  void updateBlockMaskCache(const DenseMap<VPValue *, VPValue *> &Old2New) {
    for (auto &[_, V] : BlockMaskCache) {
-      if (V == Old)
+      if (auto *New = Old2New.lookup(V)) {
+        V->replaceAllUsesWith(New);
        V = New;
+      }
    }
  }
 };

diff --git a/llvm/lib/Transforms/Vectorize/VPlanPredicator.cpp b/llvm/lib/Transforms/Vectorize/VPlanPredicator.cpp
@@ -21,11 +21,8 @@
 using namespace llvm;

 namespace {
-struct VPPredicator {
+class VPPredicator {
  using BlockMaskCacheTy = DenseMap<VPBasicBlock *, VPValue *>;
-  VPPredicator(BlockMaskCacheTy &BlockMaskCache)
-      : BlockMaskCache(BlockMaskCache) {}
-
  /// Builder to construct recipes to compute masks.
  VPBuilder Builder;

@@ -38,35 +35,45 @@ struct VPPredicator {

  BlockMaskCacheTy &BlockMaskCache;

-  /// Returns the previously computed predicate of the edge between \p Src and
-  /// \p Dst.
-  VPValue *getEdgeMask(const VPBasicBlock *Src, const VPBasicBlock *Dst) const {
-    return EdgeMaskCache.lookup({Src, Dst});
-  }
+  /// 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.
-  /// Computes and return the predicate of the edge between \p Src and \p Dst.
+  /// 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).
-  /// Computes and return the predicate of the edge between \p Src and \p Dst.
+  /// 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);

-
+
+  void setEdgeMask(const VPBasicBlock *Src, const VPBasicBlock *Dst, VPValue *Mask) {
+    EdgeMaskCache[{Src, Dst}] = Mask;
+  }
-
+
+  void setEdgeMask(const VPBasicBlock *Src, const VPBasicBlock *Dst, VPValue *Mask) {
+    EdgeMaskCache[{Src, Dst}] = Mask;
+  }
  /// Returns the *entry* mask for \p VPBB.
  VPValue *getBlockInMask(VPBasicBlock *VPBB) const {
    return BlockMaskCache.lookup(VPBB);
  }
-  }
+  }
+
-  }
+  }
+
+
  void setBlockInMask(VPBasicBlock *VPBB, VPValue *Mask) {
-  void setBlockInMask(VPBasicBlock *VPBB, VPValue *Mask) {
+  /// 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) {
-  void setBlockInMask(VPBasicBlock *VPBB, VPValue *Mask) {
+  /// 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
    // remove the need to keep a map of masks beyond the predication transform.
-    // remove the need to keep a map of masks beyond the predication transform.
+    // avoid keeping the map of masks beyond the predication transform.
-    // remove the need to keep a map of masks beyond the predication transform.
+    // avoid keeping the map of masks beyond the predication transform.
-    assert(!BlockMaskCache.contains(VPBB) && "Mask already set");
+    assert(!getBlockInMask(VPBB) && "Mask already set");
    BlockMaskCache[VPBB] = Mask;
  }

-
+
+  void setEdgeMask(VPBasicBlock *Src, VPBasicBlock *Dst, VPValue *Mask) {
+    assert(!getEdgeMask({Src, Dst}) && "Mask already set");
+    EdgeMaskCache[{Src, Dst}] = Mask;
+  }
-
+
+  void setEdgeMask(VPBasicBlock *Src, VPBasicBlock *Dst, VPValue *Mask) {
+    assert(!getEdgeMask({Src, Dst}) && "Mask already set");
+    EdgeMaskCache[{Src, Dst}] = Mask;
+  }
+  VPValue *setEdgeMask(const VPBasicBlock *Src, const VPBasicBlock *Dst,
-  VPValue *setEdgeMask(const VPBasicBlock *Src, const VPBasicBlock *Dst,
+  /// 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,
-  VPValue *setEdgeMask(const VPBasicBlock *Src, const VPBasicBlock *Dst,
+  /// 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,
+                       VPValue *Mask) {
+    assert(!getEdgeMask(Src, Dst) && "Mask already set");
+    return EdgeMaskCache[{Src, Dst}] = Mask;
+  }
+
+public:
+  VPPredicator(BlockMaskCacheTy &BlockMaskCache)
+      : BlockMaskCache(BlockMaskCache) {}
+
+  /// Returns the precomputed predicate of the edge from \p Src to \p Dst.
+  VPValue *getEdgeMask(const VPBasicBlock *Src, const VPBasicBlock *Dst) const {
+    return EdgeMaskCache.lookup({Src, Dst});
+  }
+
  /// Compute and return the mask for the vector loop header block.
  void createHeaderMask(VPBasicBlock *HeaderVPBB, bool FoldTail);

  /// Compute and return the predicate of \p VPBB, assuming that the header
  /// block of the loop is set to True or the loop mask when tail folding.
-  /// block of the loop is set to True or the loop mask when tail folding.
+  /// block of the loop is set to True, or to the loop mask when tail folding.
-  /// block of the loop is set to True or the loop mask when tail folding.
+  /// block of the loop is set to True, or to the loop mask when tail folding.
  VPValue *createBlockInMask(VPBasicBlock *VPBB);
-
-  /// Computes and return the predicate of the edge between \p Src and \p Dst.
-  VPValue *createEdgeMask(VPBasicBlock *Src, VPBasicBlock *Dst);
-
-  /// Create an edge mask for every destination of cases and/or default.
-  void createSwitchEdgeMasks(VPInstruction *SI);
 };
 } // namespace

@@ -80,40 +87,35 @@ VPValue *VPPredicator::createEdgeMask(VPBasicBlock *Src, VPBasicBlock *Dst) {

  VPValue *SrcMask = getBlockInMask(Src);

-  // The terminator has to be a branch inst!
-  if (Src->empty() || Src->getNumSuccessors() == 1) {
-    EdgeMaskCache[{Src, Dst}] = SrcMask;
-    return SrcMask;
-  }
+  // If there's a single successor, there's no terminator recipe.
+  if (Src->getNumSuccessors() == 1)
+    return setEdgeMask(Src, Dst, SrcMask);

  auto *Term = cast<VPInstruction>(Src->getTerminator());
  if (Term->getOpcode() == Instruction::Switch) {
    createSwitchEdgeMasks(Term);
    return getEdgeMask(Src, Dst);
  }

-  auto *BI = cast<VPInstruction>(Src->getTerminator());
-  assert(BI->getOpcode() == VPInstruction::BranchOnCond);
-  if (Src->getSuccessors()[0] == Src->getSuccessors()[1]) {
-    EdgeMaskCache[{Src, Dst}] = SrcMask;
-    return SrcMask;
-  }
+  assert(Term->getOpcode() == VPInstruction::BranchOnCond &&
+         "Unsupported terminator");
+  if (Src->getSuccessors()[0] == Src->getSuccessors()[1])
+    return setEdgeMask(Src, Dst, SrcMask);

-  EdgeMask = BI->getOperand(0);
+  EdgeMask = Term->getOperand(0);
  assert(EdgeMask && "No Edge Mask found for condition");

  if (Src->getSuccessors()[0] != Dst)
-    EdgeMask = Builder.createNot(EdgeMask, BI->getDebugLoc());
+    EdgeMask = Builder.createNot(EdgeMask, Term->getDebugLoc());

  if (SrcMask) { // Otherwise block in-mask is all-one, no need to AND.
    // The bitwise 'And' of SrcMask and EdgeMask introduces new UB if SrcMask
    // is false and EdgeMask is poison. Avoid that by using 'LogicalAnd'
    // instead which generates 'select i1 SrcMask, i1 EdgeMask, i1 false'.
-    EdgeMask = Builder.createLogicalAnd(SrcMask, EdgeMask, BI->getDebugLoc());
+    EdgeMask = Builder.createLogicalAnd(SrcMask, EdgeMask, Term->getDebugLoc());
  }

-  EdgeMaskCache[{Src, Dst}] = EdgeMask;
-  return EdgeMask;
+  return setEdgeMask(Src, Dst, EdgeMask);
 }

 VPValue *VPPredicator::createBlockInMask(VPBasicBlock *VPBB) {
@@ -131,7 +133,7 @@ VPValue *VPPredicator::createBlockInMask(VPBasicBlock *VPBB) {
      return EdgeMask;
    }

-    if (!BlockMask) { // BlockMask has its initialized nullptr value.
+    if (!BlockMask) { // BlockMask has its initial nullptr value.
      BlockMask = EdgeMask;
      continue;
    }
@@ -159,11 +161,9 @@ void VPPredicator::createHeaderMask(VPBasicBlock *HeaderVPBB, bool FoldTail) {
  auto *IV = new VPWidenCanonicalIVRecipe(Plan.getCanonicalIV());
  HeaderVPBB->insert(IV, NewInsertionPoint);

-  VPBuilder::InsertPointGuard Guard(Builder);
  Builder.setInsertPoint(HeaderVPBB, NewInsertionPoint);
-  VPValue *BlockMask = nullptr;
  VPValue *BTC = Plan.getOrCreateBackedgeTakenCount();
-  BlockMask = Builder.createICmp(CmpInst::ICMP_ULE, IV, BTC);
+  VPValue *BlockMask = Builder.createICmp(CmpInst::ICMP_ULE, IV, BTC);
  setBlockInMask(HeaderVPBB, BlockMask);
 }

@@ -179,7 +179,7 @@ void VPPredicator::createSwitchEdgeMasks(VPInstruction *SI) {
  for (const auto &[Idx, Succ] :
       enumerate(ArrayRef(Src->getSuccessors()).drop_front())) {
    VPBasicBlock *Dst = cast<VPBasicBlock>(Succ);
-    assert(!EdgeMaskCache.contains({Src, Dst}) && "Edge masks already created");
+    assert(!getEdgeMask(Src, Dst) && "Edge masks already created");
    //  Cases whose destination is the same as default are redundant and can
    //  be ignored - they will get there anyhow.
    if (Dst == DefaultDst)
@@ -202,7 +202,7 @@ void VPPredicator::createSwitchEdgeMasks(VPInstruction *SI) {
      Mask = Builder.createOr(Mask, V);
    if (SrcMask)
      Mask = Builder.createLogicalAnd(SrcMask, Mask);
-    EdgeMaskCache[{Src, Dst}] = Mask;
+    setEdgeMask(Src, Dst, Mask);

    // 2. Create the mask for the default destination, which is reached if
    // none of the cases with destination != default destination are taken.
@@ -216,7 +216,7 @@ void VPPredicator::createSwitchEdgeMasks(VPInstruction *SI) {
    if (SrcMask)
      DefaultMask = Builder.createLogicalAnd(SrcMask, DefaultMask);
  }
-  EdgeMaskCache[{Src, DefaultDst}] = DefaultMask;
+  setEdgeMask(Src, DefaultDst, DefaultMask);
 }

 void VPlanTransforms::predicateAndLinearize(
@@ -229,7 +229,12 @@ void VPlanTransforms::predicateAndLinearize(
  ReversePostOrderTraversal<VPBlockShallowTraversalWrapper<VPBlockBase *>> RPOT(
      Header);
  VPPredicator Predicator(BlockMaskCache);
-  for (VPBasicBlock *VPBB : VPBlockUtils::blocksOnly<VPBasicBlock>(RPOT)) {
+  for (VPBlockBase *VPB : RPOT) {
+    // Only regions with only VPBBs are supported at the moment.
-    // Only regions with only VPBBs are supported at the moment.
+    // Non-outer regions with VPBBs only are supported at the moment.
-    // Only regions with only VPBBs are supported at the moment.
+    // Non-outer regions with VPBBs only are supported at the moment.
+    auto *VPBB = cast<VPBasicBlock>(VPB);
+    // Introduce the mask for VPBB, which may introduce needed edge masks, and
+    // convert all phi recipes of VPBB to blend recipes unless VPBB is the
+    // header.
    if (VPBB == Header) {
      Predicator.createHeaderMask(Header, FoldTail);
      continue;
@@ -241,53 +246,51 @@ void VPlanTransforms::predicateAndLinearize(

    Predicator.createBlockInMask(VPBB);

-    for (VPWidenPHIRecipe *Phi : Phis) {
-      PHINode *IRPhi = cast<PHINode>(Phi->getUnderlyingValue());
-
-      unsigned NumIncoming = IRPhi->getNumIncomingValues();
-
-      // We know that all PHIs in non-header blocks are converted into selects,
+    for (VPWidenPHIRecipe *PhiR : Phis) {
+      // The non-header Phi is converted into a Blend recipe below,
      // so we don't have to worry about the insertion order and we can just use
      // the builder. At this point we generate the predication tree. There may
      // be duplications since this is a simple recursive scan, but future
      // optimizations will clean it up.

      SmallVector<VPValue *, 2> OperandsWithMask;
+      unsigned NumIncoming = PhiR->getNumIncoming();
      for (unsigned In = 0; In < NumIncoming; In++) {
-      for (unsigned In = 0; In < NumIncoming; In++) {
+      unsigned NumIncoming = PhiR->getNumIncoming();
+      for (unsigned In = 0; In < NumIncoming; In++) {
-      for (unsigned In = 0; In < NumIncoming; In++) {
+      unsigned NumIncoming = PhiR->getNumIncoming();
+      for (unsigned In = 0; In < NumIncoming; In++) {
-        const VPBasicBlock *Pred = Phi->getIncomingBlock(In);
-        OperandsWithMask.push_back(Phi->getIncomingValue(In));
+        const VPBasicBlock *Pred = PhiR->getIncomingBlock(In);
+        OperandsWithMask.push_back(PhiR->getIncomingValue(In));
        VPValue *EdgeMask = Predicator.getEdgeMask(Pred, VPBB);
        if (!EdgeMask) {
          assert(In == 0 && "Both null and non-null edge masks found");
-          assert(all_equal(Phi->operands()) &&
+          assert(all_equal(PhiR->operands()) &&
                 "Distinct incoming values with one having a full mask");
          break;
        }
        OperandsWithMask.push_back(EdgeMask);
      }
+      PHINode *IRPhi = cast<PHINode>(PhiR->getUnderlyingValue());
      auto *Blend = new VPBlendRecipe(IRPhi, OperandsWithMask);
-      auto *Blend = new VPBlendRecipe(IRPhi, OperandsWithMask);
+      PHINode *IRPhi = cast<PHINode>(Phi->getUnderlyingValue());
+      auto *Blend = new VPBlendRecipe(IRPhi, OperandsWithMask);
-      auto *Blend = new VPBlendRecipe(IRPhi, OperandsWithMask);
+      PHINode *IRPhi = cast<PHINode>(Phi->getUnderlyingValue());
+      auto *Blend = new VPBlendRecipe(IRPhi, OperandsWithMask);
-      Blend->insertBefore(Phi);
-      Phi->replaceAllUsesWith(Blend);
-      Phi->eraseFromParent();
+      Blend->insertBefore(PhiR);
+      PhiR->replaceAllUsesWith(Blend);
+      PhiR->eraseFromParent();
    }
  }

+  // Linearize the blocks of the loop into one serial chain.
  VPBlockBase *PrevVPBB = nullptr;
-  VPBlockBase *PrevVPBB = nullptr;
+  // Linearize the blocks of the loop into one serial chain.
+  VPBlockBase *PrevVPBB = nullptr;
-  VPBlockBase *PrevVPBB = nullptr;
+  // Linearize the blocks of the loop into one serial chain.
+  VPBlockBase *PrevVPBB = nullptr;
  for (VPBasicBlock *VPBB : VPBlockUtils::blocksOnly<VPBasicBlock>(RPOT)) {
    // Handle VPBBs down to the latch.
-    if (PrevVPBB && VPBB == LoopRegion->getExiting()) {
-      VPBlockUtils::connectBlocks(PrevVPBB, VPBB);
+    if (VPBB == LoopRegion->getExiting()) {
+      if (PrevVPBB)
+        VPBlockUtils::connectBlocks(PrevVPBB, VPBB);
      break;
    }

    auto Successors = to_vector(VPBB->getSuccessors());
    if (Successors.size() > 1)
      VPBB->getTerminator()->eraseFromParent();

-    // Flatten the CFG in the loop. Masks for blocks have already been
-    // generated and added to recipes as needed. To do so, first disconnect
-    // VPBB from its successors. Then connect VPBB to the previously visited
-    // VPBB.
+    // Flatten the CFG in the loop. To do so, first disconnect VPBB from its
+    // successors. Then connect VPBB to the previously visited VPBB.
    for (auto *Succ : Successors)
      VPBlockUtils::disconnectBlocks(VPBB, Succ);
    if (PrevVPBB)

diff --git a/llvm/lib/Transforms/Vectorize/VPlanTransforms.h b/llvm/lib/Transforms/Vectorize/VPlanTransforms.h
@@ -223,12 +223,12 @@ struct VPlanTransforms {
  static void narrowInterleaveGroups(VPlan &Plan, ElementCount VF,
                                     unsigned VectorRegWidth);

-  /// Predicate and linearize the control-flow in the top-level loop region of
+  /// Predicate and linearize the control-flow in the only loop region of
  /// \p Plan. If \p FoldTail is true, also create a mask guarding the loop
-  /// \p Plan. If \p FoldTail is true, also create a mask guarding the loop
+  /// \p Plan. If \p FoldTail is true, create a mask guarding the loop
-  /// \p Plan. If \p FoldTail is true, also create a mask guarding the loop
+  /// \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 \p BlockMaskCache, which in turn is temporarily used for wide
-  /// recipe construction. This argument is temporary and will be removed in the
-  /// future.
+  /// added to \p BlockMaskCache, which in turn will temporarily be used later
-  /// added to \p BlockMaskCache, which in turn will temporarily be used later
+  /// added to \p BlockMaskCache in order to be used later
-  /// added to \p BlockMaskCache, which in turn will temporarily be used later
+  /// added to \p BlockMaskCache in order to be used later
+  /// for wide recipe construction. This argument is temporary and will be
+  /// removed in the future.
  static void
  predicateAndLinearize(VPlan &Plan, bool FoldTail,
                        DenseMap<VPBasicBlock *, VPValue *> &BlockMaskCache);
-Original file line number
+Diff line change
@@ -209,10 +209,12 @@ class VPRecipeBuilder {
         return Plan.getOrAddLiveIn(V);
       }
-      void updateBlockMaskCache(VPValue *Old, VPValue *New) {
+      void updateBlockMaskCache(const DenseMap<VPValue *, VPValue *> &Old2New) {
         for (auto &[_, V] : BlockMaskCache) {
-          if (V == Old)
+          if (auto *New = Old2New.lookup(V)) {
+            V->replaceAllUsesWith(New);
             Copy link

  
      
    
  

  
      

  
    Collaborator


      

  

  
    
      

      
            ayalz
  

      

      

      


        May 18, 2025


      
    

  


        
      
  
  
  
    

    There was a problem hiding this comment.


  

 
  
    

    Choose a reason for hiding this comment

    
      The reason will be displayed to describe this comment to others. Learn more.
    

    
      
      


  


  
    
      nit: worth removing V from Old2New now?
    
  
  


    

        
      
  
  
    
  
  
  
    
    Sorry, something went wrong.
  

  
    
  
    
      

              Uh oh!

              
There was an error while loading. Please reload this page.


  
  


          
      
  
    
    
      
        
            
    All reactions
  


          
          
        
      
    

    



  
        
  
    
        
    
  


      
          
  
      
            Copy link

  
      
    
  

  
      

  
    Contributor


      

  Author


  

  
    
      

      
            fhahn
  

      

      

      


        May 21, 2025


      
    

  


        
      
  
  
  
    

    There was a problem hiding this comment.


  

 
  
    

    Choose a reason for hiding this comment

    
      The reason will be displayed to describe this comment to others. Learn more.
    

    
      
      


  


  
    
      Cannot be done for now, as Old2New is used to erase old recipes after updateBlockMaskCache
    
  
  


    

        
      
  
  
    
  
  
  
    
    Sorry, something went wrong.
  

  
    
  
    
      

              Uh oh!

              
There was an error while loading. Please reload this page.


  
  


          
      
  
    
    
      
        
            
    All reactions
             V = New;
+          }
         }
       }
     };
-          Expand Down