[VPlan] Proactively create mask for tail-folding up-front (NFCI).

[fhahn]

Summary

Split off mask creation for tail folding and proactively create the mask for
the header block.

This simplifies createBlockInMask.

[Ayal]

Split off mask creation for tail folding and always create the mask for
the header block.

"always" >> "proactively"?

Down the road tail folding should be considered as a VPlan2VPlan transformation along with masking/predication.

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

8011

nit (independent of this patch): more accurately prefixed getOrCreate.

8024–8025

nit: indent

The comment about "Loop incoming mask" is now obsolete, it refers to BB being the header.

8024–8025

add assert message; a header mask is handled earlier by createTailFoldHeaderMask(), and now must have been found in cache.

8042

Have it return void?

Note (independent of this patch): this is truly create, except for useActiveLaneMaskForControlFlow case...

8046

assert that we create it only once?

8050

nit: here a null BlockMask will not be used. Better define it later closer to where it's actually set and used.

8053

nit: best assert earlier, first thing.

8062

nit (independent of this patch): some explanation of using active lane mask (below, not for control flow - above) which does use IV and TC, is missing here.

8776–8779

Worth a comment to explain that the mask of the header block is created here proactively, if needed - when folding tail, leaving the masks of other blocks to be created on-demand, using the header block mask.

8835–8836

Why/Is this change needed/related to tail folding? Any (header) phi recipe should indeed be placed along with other phi's rather than be appended after non-phi recipes. But are there now other non-phi, non-trunc Instr's that generate non-widen-int-or-fp-induction header phi recipes?

If needed, above comment needs update.

Perhaps better optimize truncating IV's later as a VPlan2VPlan transformation, as in optimizeInductions().

llvm/lib/Transforms/Vectorize/VPRecipeBuilder.h

139–140

(nit: this assumption will now hold ;-)

[fhahn]

Address latest comments, thanks!

Down the road tail folding should be considered as a VPlan2VPlan transformation along with masking/predication.

Agreed! This would effectively require adjusting all existing masks with the loop mask.

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

8024–8025

Fixed indent and moved this code up, also more accurately checking for BB == OrigLoop->getHeader() which should be the only relevant case.

8024–8025

Assert removed, as we now handle OrigLoop->getHeader() == BB directly.

8042

Updated to return void

8046

This isn't needed as createTailFoldHeaderMask is called once and nothing will be cached then. Remove the code.

8050

Moved, also dropped the comment; there always will be a non-all-one mask.

8053

Moved, thanks!

8776–8779

Thanks, added a comment!

8835–8836

I updated the comment for now. Doing as a VPlan2VPlan transform should be feasible in general, but it would require TTI to check if the truncate would be free.

llvm/lib/Transforms/Vectorize/VPRecipeBuilder.h

139–140

Clarified comment that the header mask will either be set to true or the loop mask when tail-folding.

[fhahn]

Rebase and add missing early return. Ping :)

[fhahn]

I shared a patch to perform tail-folding as VPlan2VPlan transform in D157713. It requires a bit more work to get the right APIs and things into place, but separating out the header mask creation code as in the current patch should be good first step in that direction.

[ABataev]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

8016

It may create a new item in BlockMaskCache with nullptr value. Do you need to check that the user actually expects it?

[Ayal]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

8016

Agreed. Better explicitly always set BlockMaskCache[] for header BB, and assert BB is not the header after looking for its mask in the cache below.

8023

nit: setting BlockMask to null here mostly represents "uninitialized" - to be overwritten by the first edge mask; and to serve as the (all-true) mask for blocks free of any predecessors.

8062

nit (independent of this patch).

8777–8780

and have createHeaderMask() set the mask to null if foldTail is false?

8835–8844

nit (independent of this patch): can the check that first non phi isn't end() be dropped? I.e., inserting before end is fine and equivalent to append.
nit (given longer explanation): consider dealing with the simpler appending case first.

8845–8847

ahh, thanks for the clarification!

8847

nit: "before" here means both in time and space. Perhaps "non-phi recipes setting header mask are introduced earlier than regular header phi recipes, and should appear after them," would be clearer. (Related to all header phi's, not only inductions).

[fhahn]

Address latest comments, thanks!

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

8016

Thanks, updated to use .find() and added an assert. Clients already should handle the case the mask is nullptr, as this is used to indicate no-mask (all-true mask) in the existing code.

8062

Done in c34d049706e69cb61d05ce21ded2ac7ecef9f45f, thanks!

8777–8780

Updated, thanks!

8835–8844

I dropped the check from the condition in 57a6f6579c97 and moved it to the assert, thanks!

8845–8847

Adjusted, thanks!

8847

Reworded, thanks!

[Ayal]

This looks good to me, adding some thoughts that could be addresses now or later.

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

8014–8024

?

8053

Would it be better to have one switch (TFStyle) whose cases compute BlockMask, followed by placing it in the cache, where None refers to not folding the tail?

8076

This else case of not using active lane mask is aka DataWithoutLaneMask.

ActiveLaneMask is semantically equivalent to an ICmpULE (except for wrapping TC?), so down the road we could start canonically with one and if desired replace it with the other, or with an ActiveLaneMaskPhi, or None, as VPlan2VPlan transforms.

8773

Sh/Could addCanonicalIVRecipes() move to RecipeBuilder?

8835–8836

Would something like the following be better?

auto InsertionPoint = isa<VPHeaderPHIRecipe>(Recipe) ? VPBB->getFirstNonPhi() : VBPP->end();
Recipe->insertBefore(*VPBB, InsertionPoint);

with an

assert ((!isa<VPHeaderPHIRecipe>(Recipe) ||
         VPBB->getFirstNonPhi() == VPBB->end() ||
         CM.foldTailByMasking() || isa<TruncInst>(Instr)) &&
               "unexpected recipe needs moving");

?

Note that phi recipes in VPBB's other than header VPBB also need to be placed in the phi section of their VPBB, in general. However, (a) this currently includes only VPPredInstPHIRecipe, which is introduced into VPlan only later, and (b) all current cases of out-of-place recipes occur in the header only - its mask and a trunc of an induction phi.

llvm/lib/Transforms/Vectorize/VPRecipeBuilder.h

146

nit: seems slightly more logical to place createHeaderMask() first, before createBlockInMask(), which in turn continues to be followed by createEdgeMask.

[fhahn]

Address latest comments, thanks! I am planning on landing this soon.

[fhahn]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

8014–8024

Simplified as suggested, thanks!

8053

I'll share a patch soon to add active-lane-mask as optimization.

8076

I'll share a patch soon to add active-lane-mask as optimization.

8773

Probably, let me check.

8835–8836

Sounds good but I think it would be better to do that separately from the current change.

llvm/lib/Transforms/Vectorize/VPRecipeBuilder.h

146

Yep, reordered.

[Ayal]

Couple of minor nits.

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

8042

nit: reorder implementations of createHeaderMask() and createBlockInMask(), matching the reordering of their declarations.

8777

[fhahn]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

8042

Fixed in the committed version, thanks!

8777

Fixed in the committed version, thanks!

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

VPValue *False = Plan.getVPValueOrAddLiveIn(
         ConstantInt::getFalse(BI->getCondition()->getType()));
     EdgeMask =
         Builder.createSelect(SrcMask, EdgeMask, False, BI->getDebugLoc());
   }
 
   return EdgeMaskCache[Edge] = EdgeMask;
 }
 
 VPValue *VPRecipeBuilder::createBlockInMask(BasicBlock *BB, VPlan &Plan) {
   assert(OrigLoop->contains(BB) && "Block is not a part of a loop");
 
   // Look for cached value.
   BlockMaskCacheTy::iterator BCEntryIt = BlockMaskCache.find(BB);
   if (BCEntryIt != BlockMaskCache.end())
     return BCEntryIt->second;
 
+  assert(OrigLoop->getHeader() != BB &&
+         "Loop header must have cached block mask");
+
   // All-one mask is modelled as no-mask following the convention for masked
   // load/store/gather/scatter. Initialize BlockMask to no-mask.
   VPValue *BlockMask = nullptr;
+  // This is the block mask. We OR all incoming edges.
+  for (auto *Predecessor : predecessors(BB)) {
+    VPValue *EdgeMask = createEdgeMask(Predecessor, BB, Plan);
+    if (!EdgeMask) // Mask of predecessor is all-one so mask of block is too.
+      return BlockMaskCache[BB] = EdgeMask;
 
-  if (OrigLoop->getHeader() == BB) {
-    if (!CM.blockNeedsPredicationForAnyReason(BB))
-      return BlockMaskCache[BB] = BlockMask; // Loop incoming mask is all-one.
+    if (!BlockMask) { // BlockMask has its initialized nullptr value.
+      BlockMask = EdgeMask;
+      continue;
+    }
 
-    assert(CM.foldTailByMasking() && "must fold the tail");
+    BlockMask = Builder.createOr(BlockMask, EdgeMask, {});
+  }
+
+  return BlockMaskCache[BB] = BlockMask;
+}
+
+void VPRecipeBuilder::createHeaderMask(VPlan &Plan) {
+  BasicBlock *Header = OrigLoop->getHeader();
+
+  // When not folding the tail, use nullptr to model all-true mask.
+  if (!CM.foldTailByMasking()) {
+    BlockMaskCache[Header] = nullptr;
+    return;
+  }
 
-    // If we're using the active lane mask for control flow, then we get the
-    // mask from the active lane mask PHI that is cached in the VPlan.
-    TailFoldingStyle TFStyle = CM.getTailFoldingStyle();
-    if (useActiveLaneMaskForControlFlow(TFStyle))
-      return BlockMaskCache[BB] = Plan.getActiveLaneMaskPhi();
+  if (useActiveLaneMaskForControlFlow(TFStyle)) {
+    BlockMaskCache[Header] = Plan.getActiveLaneMaskPhi();
+    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
-    // constructing the desired canonical IV in the header block as its first
-    // non-phi instructions.
 
-    VPBasicBlock *HeaderVPBB = Plan.getVectorLoopRegion()->getEntryBasicBlock();
-    auto NewInsertionPoint = HeaderVPBB->getFirstNonPhi();
-    auto *IV = new VPWidenCanonicalIVRecipe(Plan.getCanonicalIV());
-    HeaderVPBB->insert(IV, NewInsertionPoint);
 
-    VPBuilder::InsertPointGuard Guard(Builder);
-    Builder.setInsertPoint(HeaderVPBB, NewInsertionPoint);
-    if (useActiveLaneMask(TFStyle)) {
-      VPValue *TC = Plan.getTripCount();
-      BlockMask = Builder.createNaryOp(VPInstruction::ActiveLaneMask, {IV, TC},
-                                       nullptr, "active.lane.mask");
-    } else {
-      VPValue *BTC = Plan.getOrCreateBackedgeTakenCount();
-      BlockMask = Builder.createNaryOp(VPInstruction::ICmpULE, {IV, BTC});
-    }
-    return BlockMaskCache[BB] = BlockMask;
-  }
-
-  // This is the block mask. We OR all incoming edges.
-  for (auto *Predecessor : predecessors(BB)) {
-    VPValue *EdgeMask = createEdgeMask(Predecessor, BB, Plan);
-    if (!EdgeMask) // Mask of predecessor is all-one so mask of block is too.
-      return BlockMaskCache[BB] = EdgeMask;
-
-    if (!BlockMask) { // BlockMask has its initialized nullptr value.
-      BlockMask = EdgeMask;
-      continue;
-    }
-
-    BlockMask = Builder.createOr(BlockMask, EdgeMask, {});
-  }
-
-  return BlockMaskCache[BB] = BlockMask;
+  VPValue *BlockMask = nullptr;
+  BlockMaskCache[Header] = BlockMask;
 }
 
 VPRecipeBase *VPRecipeBuilder::tryToWidenMemory(Instruction *I,
                                                 ArrayRef<VPValue *> Operands,
                                                 VFRange &Range,
                                                 VPlanPtr &Plan) {
   assert((isa<LoadInst>(I) || isa<StoreInst>(I)) &&
          "Must be called with either a load or store");
LoopVectorizationPlanner::tryToBuildVPlanWithVPRecipes(VFRange &Range) {
   // so this function is better to be conservative, rather than to split
   // it up into different VPlans.
   bool IVUpdateMayOverflow = false;
   for (ElementCount VF : Range)
     IVUpdateMayOverflow |= !isIndvarOverflowCheckKnownFalse(&CM, VF);
 
   Instruction *DLInst =
       getDebugLocFromInstOrOperands(Legal->getPrimaryInduction());
   addCanonicalIVRecipes(*Plan, Legal->getWidestInductionType(),
                         DLInst ? DLInst->getDebugLoc() : DebugLoc(),
                         CM.getTailFoldingStyle(IVUpdateMayOverflow));
 
+  // Proactively create header. Masks for other blocks are created on demand.
+  RecipeBuilder.createHeaderMask(*Plan);
+
   // Scan the body of the loop in a topological order to visit each basic block
   // after having visited its predecessor basic blocks.
   LoopBlocksDFS DFS(OrigLoop);
   DFS.perform(LI);
 
   VPBasicBlock *VPBB = HeaderVPBB;
   for (BasicBlock *BB : make_range(DFS.beginRPO(), DFS.endRPO())) {
     // Relevant instructions from basic block BB will be grouped into VPRecipe
     // ingredients and fill a new VPBasicBlock.
for (Instruction &I : drop_end(BB->instructionsWithoutDebug(false))) {
       }
       // Otherwise, add the new recipe.
       VPRecipeBase *Recipe = cast<VPRecipeBase *>(RecipeOrValue);
       for (auto *Def : Recipe->definedValues()) {
         auto *UV = Def->getUnderlyingValue();
         Plan->addVPValue(UV, Def);
       }
 
       RecipeBuilder.setRecipe(Instr, Recipe);
-      if (isa<VPWidenIntOrFpInductionRecipe>(Recipe)) {
-        // VPWidenIntOrFpInductionRecipes must be kept in the phi section of
-        // HeaderVPBB. VPWidenIntOrFpInductionRecipes for optimized truncates
-        // may be generated after non-phi recipes and need to be moved to the
-        // phi section of HeaderVPBB.
+      if (isa<VPHeaderPHIRecipe>(Recipe)) {
+        // VPHeaderPHIRecipes must be kept in the phi section of HeaderVPBB. In
+        // the following cases, VPHeaderPHIRecipes may be created after non-phi
+        // recipes and need to be moved to the phi section of HeaderVPBB:
+        // * tail-folding (non-phi recipes computing the header mask are
+        // introduced earlier than regular header phi recipes, and should appear
+        // after them)
+        // * Optimizing truncates to VPWidenIntOrFpInductionRecipe.
+
         assert((HeaderVPBB->getFirstNonPhi() == VPBB->end() ||
-                isa<TruncInst>(Instr)) &&
+                CM.foldTailByMasking() || isa<TruncInst>(Instr)) &&
                "unexpected recipe needs moving");
         Recipe->insertBefore(*HeaderVPBB, HeaderVPBB->getFirstNonPhi());
       } else
         VPBB->appendRecipe(Recipe);
     }
 
     VPBlockUtils::insertBlockAfter(new VPBasicBlock(), VPBB);
     VPBB = cast<VPBasicBlock>(VPBB->getSingleSuccessor());
   }

llvm/lib/Transforms/Vectorize/VPRecipeBuilder.h

public:
   void setRecipe(Instruction *I, VPRecipeBase *R) {
     if (!Ingredient2Recipe.count(I))
       return;
     assert(Ingredient2Recipe[I] == nullptr &&
            "Recipe already set for ingredient");
     Ingredient2Recipe[I] = R;
   }
 
+  /// Create the mask for the vector loop header block.
+  void createHeaderMask(VPlan &Plan);
+
   /// A helper function that computes the predicate of the block BB, assuming
-  /// that the header block of the loop is set to True. It returns the *entry*
-  /// mask for the block BB.
+  /// that the header block of the loop is set to True or the loop mask when
+  /// tail folding. It returns the *entry* mask for the block BB.
   VPValue *createBlockInMask(BasicBlock *BB, VPlan &Plan);
 
   /// A helper function that computes the predicate of the edge between SRC
   /// and DST.
   VPValue *createEdgeMask(BasicBlock *Src, BasicBlock *Dst, VPlan &Plan);
 
   /// Mark given ingredient for recording its recipe once one is created for
   /// it.
   void recordRecipeOf(Instruction *I) {
     assert((!Ingredient2Recipe.count(I) || Ingredient2Recipe[I] == nullptr) &&
            "Recipe already set for ingredient");
     Ingredient2Recipe[I] = nullptr;
   }