[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

8148

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

8161

nit: indent

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

8162

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

8181

Have it return void?

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

8185

assert that we create it only once?

8189

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

8192

nit: best assert earlier, first thing.

8201

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.

8916

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.

8984

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

137

(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

8161

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

8162

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

8181

Updated to return void

8185

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

8189

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

8192

Moved, thanks!

8916

Thanks, added a comment!

8984

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

137

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

8153

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

8153

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.

8157

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.

8201

nit (independent of this patch).

8917–8920

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

8983–8985

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.

8986

ahh, thanks for the clarification!

8988

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

8153

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.

8201

Done in c34d049706e69cb61d05ce21ded2ac7ecef9f45f, thanks!

8917–8920

Updated, thanks!

8983–8985

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

8986

Adjusted, thanks!

8988

Reworded, thanks!

[Ayal]

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

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

8151–8161

?

8192

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?

8215

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.

8913

Sh/Could addCanonicalIVRecipes() move to RecipeBuilder?

8981–8984

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

143

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

8151–8161

Simplified as suggested, thanks!

8192

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

8215

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

8913

Probably, let me check.

8981–8984

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

llvm/lib/Transforms/Vectorize/VPRecipeBuilder.h

143

Yep, reordered.

[Ayal]

Couple of minor nits.

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

8181

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

8917

[fhahn]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

8181

Fixed in the committed version, thanks!

8917

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;
 
   // 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;
 
-  if (OrigLoop->getHeader() == BB) {
     if (!CM.blockNeedsPredicationForAnyReason(BB))
       return BlockMaskCache[BB] = BlockMask; // Loop incoming mask is all-one.
+  assert (OrigLoop->getHeader() != BB);
+
+  // 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 *VPRecipeBuilder::createTailFoldHeaderMask(VPlan &Plan) {
+  BasicBlock *Header = OrigLoop->getHeader();
+  // Look for cached value.
+  BlockMaskCacheTy::iterator BCEntryIt = BlockMaskCache.find(Header);
+  if (BCEntryIt != BlockMaskCache.end())
+    return BCEntryIt->second;
+
+  // 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;
 
-    assert(CM.foldTailByMasking() && "must fold the tail");
 
-    // 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();
+    return BlockMaskCache[Header] = Plan.getActiveLaneMaskPhi();
 
-    // 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, HeaderVPBB->getFirstNonPhi());
 
-    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;
+  return 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));
 
+  if (CM.foldTailByMasking())
+    RecipeBuilder.createTailFoldHeaderMask(*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 : 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) &&
+      if (isa<VPHeaderPHIRecipe>(Recipe) &&
           HeaderVPBB->getFirstNonPhi() != VPBB->end()) {
         // Move VPWidenIntOrFpInductionRecipes for optimized truncates to the
         // phi section of HeaderVPBB.
-        assert(isa<TruncInst>(Instr));
+        assert(CM.foldTailByMasking() || isa<TruncInst>(Instr));
         Recipe->insertBefore(*HeaderVPBB, HeaderVPBB->getFirstNonPhi());
       } else
         VPBB->appendRecipe(Recipe);
     }
 
     VPBlockUtils::insertBlockAfter(new VPBasicBlock(), VPBB);
     VPBB = cast<VPBasicBlock>(VPBB->getSingleSuccessor());
   }
 
   // After here, VPBB should not be used.

llvm/lib/Transforms/Vectorize/VPRecipeBuilder.h

void setRecipe(Instruction *I, VPRecipeBase *R) {
     if (!Ingredient2Recipe.count(I))
       return;
     assert(Ingredient2Recipe[I] == nullptr &&
            "Recipe already set for ingredient");
     Ingredient2Recipe[I] = R;
   }
 
   /// 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.
   VPValue *createBlockInMask(BasicBlock *BB, VPlan &Plan);
 
+  /// Create the mask for the vector loop header block when tail-folding.
+  VPValue *createTailFoldHeaderMask(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) &&