[VPlan, VP] Introduce new live-in values and refactoring to support EVL recipe

[vkmr]

Summary

This patch introduces TripCount and RuntiVF as live-in VPValues whose underlying IR Values are available only at plan execution time.
Canonical induction variable is cached so it can be reused when creating EVL recipe.

This patch is second in the series of patches to enable predicated vectorization support in LV introduced in RFC D99750.

[vkmr]

Ping!

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

if (SrcMask) { // Otherwise block in-mask is all-one, no need to AND.
     VPValue *False = Plan->getOrAddVPValue(
         ConstantInt::getFalse(BI->getCondition()->getType()));
     EdgeMask = Builder.createSelect(SrcMask, EdgeMask, False);
   }
 
   return EdgeMaskCache[Edge] = EdgeMask;
 }
 
+VPValue *VPRecipeBuilder::getOrCreateIV(VPBasicBlock *VPBB, VPlanPtr &Plan) {
+  IVCacheTy::iterator IVEntryIt = IVCache.find(VPBB);
+  if (IVEntryIt != IVCache.end())
+    return IVEntryIt->second;
+
+  VPValue *IV = nullptr;
+  if (Legal->getPrimaryInduction())
+    IV = Plan->getOrAddVPValue(Legal->getPrimaryInduction());
+  else {
+    auto *IVRecipe = new VPWidenCanonicalIVRecipe();
+    Builder.getInsertBlock()->insert(IVRecipe, Builder.getInsertPoint());
+    IV = IVRecipe->getVPSingleValue();
+  }
+  return IVCache[VPBB] = IV;
+}
+
 VPValue *VPRecipeBuilder::createBlockInMask(BasicBlock *BB, VPlanPtr &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.blockNeedsPredication(BB))
       return BlockMaskCache[BB] = BlockMask; // Loop incoming mask is all-one.
 
     // Create the block in mask as the first non-phi instruction in the block.
     VPBuilder::InsertPointGuard Guard(Builder);
     auto NewInsertionPoint = Builder.getInsertBlock()->getFirstNonPhi();
     Builder.setInsertPoint(Builder.getInsertBlock(), NewInsertionPoint);
 
     // 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.
-    VPValue *IV = nullptr;
-    if (Legal->getPrimaryInduction())
-      IV = Plan->getOrAddVPValue(Legal->getPrimaryInduction());
-    else {
-      auto IVRecipe = new VPWidenCanonicalIVRecipe();
-      Builder.getInsertBlock()->insert(IVRecipe, NewInsertionPoint);
-      IV = IVRecipe->getVPSingleValue();
-    }
+    VPValue *IV = getOrCreateIV(Builder.getInsertBlock(), Plan);
     VPValue *BTC = Plan->getOrCreateBackedgeTakenCount();
     bool TailFolded = !CM.isScalarEpilogueAllowed();
 
     if (TailFolded && CM.TTI.emitGetActiveLaneMask()) {
       // While ActiveLaneMask is a binary op that consumes the loop tripcount
       // as a second argument, we only pass the IV here and extract the
       // tripcount from the transform state where codegen of the VP instructions
       // happen.

llvm/lib/Transforms/Vectorize/VPRecipeBuilder.h

class VPRecipeBuilder {
   /// 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;
 
+  /// Hold a mapping of Basic block to the canonical vector induction VPValue
+  /// inserted for that block or the primary induction if it exists.
+  using IVCacheTy = DenseMap<VPBasicBlock *, VPValue *>;
+  IVCacheTy IVCache;
+
   // VPlan-VPlan transformations support: Hold a mapping from ingredients to
   // their recipe. To save on memory, only do so for selected ingredients,
   // marked by having a nullptr entry in this map.
   DenseMap<Instruction *, VPRecipeBase *> Ingredient2Recipe;
 
   /// Cross-iteration reduction phis for which we need to add the incoming value
   /// from the backedge after all recipes have been created.
   SmallVector<VPWidenPHIRecipe *, 4> PhisToFix;
class VPRecipeBuilder {
   /// Check if \p I has an opcode that can be widened and return a VPWidenRecipe
   /// if it can. The function should only be called if the cost-model indicates
   /// that widening should be performed.
   VPWidenRecipe *tryToWiden(Instruction *I, ArrayRef<VPValue *> Operands) const;
 
   /// Return a VPRecipeOrValueTy with VPRecipeBase * being set. This can be used to force the use as VPRecipeBase* for recipe sub-types that also inherit from VPValue.
   VPRecipeOrVPValueTy toVPRecipeResult(VPRecipeBase *R) const { return R; }
 
+  /// Insert and Cache Induction Variable
+  VPValue *getOrCreateIV(VPBasicBlock *VPBB, VPlanPtr &Plan);
+
 public:
   VPRecipeBuilder(Loop *OrigLoop, const TargetLibraryInfo *TLI,
                   LoopVectorizationLegality *Legal,
                   LoopVectorizationCostModel &CM,
                   PredicatedScalarEvolution &PSE, VPBuilder &Builder)
       : OrigLoop(OrigLoop), TLI(TLI), Legal(Legal), CM(CM), PSE(PSE),
         Builder(Builder) {}
 

llvm/lib/Transforms/Vectorize/VPlan.h

class VPlan {
   // underlying IR that will be used to implement its structural comparison
   // (operators '==' and '<').
   SetVector<VPValue *> VPExternalDefs;
 
   /// Represents the backedge taken count of the original loop, for folding
   /// the tail.
   VPValue *BackedgeTakenCount = nullptr;
 
+  /// Represents the trip count of the original loop, for computing EVL.
+  VPValue *TripCount = nullptr;
+
+  /// Represents the runtime VF. Some recipes like Vector Predicated recipes may
+  /// use runtime VF as an operand. At the time of plan construction while it is
+  /// known that this value is a loop invariant, but the corresponding IR value
+  /// is only available at plan execution once the final VF and corresponding
+  /// plan are chosen.
+  VPValue *RuntimeVF = nullptr;
+
   /// Holds a mapping between Values and their corresponding VPValue inside
   /// VPlan.
   Value2VPValueTy Value2VPValue;
 
   /// Contains all VPValues that been allocated by addVPValue directly and need
   /// to be free when the plan's destructor is called.
   SmallVector<VPValue *, 16> VPValuesToFree;
 
if (Entry) {
         Block->dropAllReferences(&DummyValue);
 
       VPBlockBase::deleteCFG(Entry);
     }
     for (VPValue *VPV : VPValuesToFree)
       delete VPV;
     if (BackedgeTakenCount)
       delete BackedgeTakenCount;
+    if (TripCount)
+      delete TripCount;
+    if (RuntimeVF)
+      delete RuntimeVF;
     for (VPValue *Def : VPExternalDefs)
       delete Def;
   }
 
   /// Generate the IR code for this VPlan.
   void execute(struct VPTransformState *State);
 
   VPBlockBase *getEntry() { return Entry; }
   const VPBlockBase *getEntry() const { return Entry; }
 
   VPBlockBase *setEntry(VPBlockBase *Block) {
     Entry = Block;
     Block->setPlan(this);
     return Entry;
   }
 
   /// The backedge taken count of the original loop.
   VPValue *getOrCreateBackedgeTakenCount() {
     if (!BackedgeTakenCount)
       BackedgeTakenCount = new VPValue();
     return BackedgeTakenCount;
   }
 
+  /// The trip count of the original loop.
+  VPValue *getOrCreateTripCount() {
+    if (!TripCount)
+      TripCount = new VPValue();
+    return TripCount;
+  }
+
+  /// A VPValue representing the loop invariant runtime VF to be expanded at
+  /// paln execution.
+  VPValue *getOrCreateRuntimeVF() {
+    if (!RuntimeVF)
+      RuntimeVF = new VPValue();
+    return RuntimeVF;
+  }
+
   void addVF(ElementCount VF) { VFs.insert(VF); }
 
   bool hasVF(ElementCount VF) { return VFs.count(VF); }
 
   const std::string &getName() const { return Name; }
 
   void setName(const Twine &newName) { Name = newName.str(); }
 

llvm/lib/Transforms/Vectorize/VPlan.cpp

void VPInstruction::print(raw_ostream &O, const Twine &Indent,
   }
 }
 #endif
 
 /// Generate the code inside the body of the vectorized loop. Assumes a single
 /// LoopVectorBody basic-block was created for this. Introduce additional
 /// basic-blocks as needed, and fill them all.
 void VPlan::execute(VPTransformState *State) {
+  IRBuilder<> Builder(State->CFG.PrevBB->getTerminator());
+
+  // -3 Check if the trip count is needed, if so build it.
+  if (TripCount && TripCount->getNumUsers()) {
+    Value *TC = State->TripCount;
+    for (unsigned Part = 0, UF = State->UF; Part < UF; ++Part)
+      State->set(TripCount, TC, Part);
+  }
+
+  // -2 Set the runtime VF if it is needed.
+  if (RuntimeVF && RuntimeVF->getNumUsers()) {
+    Value *RuntimeVFVal =
+        getRuntimeVF(Builder, Builder.getInt32Ty(), State->VF);
+    for (unsigned Part = 0, UF = State->UF; Part < UF; ++Part)
+      State->set(RuntimeVF, RuntimeVFVal, Part);
+  }
+
   // -1. Check if the backedge taken count is needed, and if so build it.
   if (BackedgeTakenCount && BackedgeTakenCount->getNumUsers()) {
     Value *TC = State->TripCount;
-    IRBuilder<> Builder(State->CFG.PrevBB->getTerminator());
     auto *TCMO = Builder.CreateSub(TC, ConstantInt::get(TC->getType(), 1),
                                    "trip.count.minus.1");
     auto VF = State->VF;
     Value *VTCMO =
         VF.isScalar() ? TCMO : Builder.CreateVectorSplat(VF, TCMO, "broadcast");
     for (unsigned Part = 0, UF = State->UF; Part < UF; ++Part)
       State->set(BackedgeTakenCount, VTCMO, Part);
   }
void VPlanPrinter::dump() {
   OS << "graph [labelloc=t, fontsize=30; label=\"Vectorization Plan";
   if (!Plan.getName().empty())
     OS << "\\n" << DOT::EscapeString(Plan.getName());
   if (Plan.BackedgeTakenCount) {
     OS << ", where:\\n";
     Plan.BackedgeTakenCount->print(OS, SlotTracker);
     OS << " := BackedgeTakenCount";
   }
+  if (Plan.TripCount) {
+    OS << "\\n";
+    Plan.RuntimeVF->print(OS, SlotTracker);
+    OS << " := TripCount";
+  }
+  if (Plan.RuntimeVF) {
+    OS << "\\n";
+    Plan.RuntimeVF->print(OS, SlotTracker);
+    OS << " := RuntimeVF";
+  }
   OS << "\"]\n";
   OS << "node [shape=rect, fontname=Courier, fontsize=30]\n";
   OS << "edge [fontname=Courier, fontsize=30]\n";
   OS << "compound=true\n";
 
   for (const VPBlockBase *Block : depth_first(Plan.getEntry()))
     dumpBlock(Block);
 

 void VPSlotTracker::assignSlots(const VPlan &Plan) {
 
   for (const VPValue *V : Plan.VPExternalDefs)
     assignSlot(V);
 
   if (Plan.BackedgeTakenCount)
     assignSlot(Plan.BackedgeTakenCount);
 
+  if (Plan.TripCount)
+    assignSlot(Plan.TripCount);
+
+  if (Plan.RuntimeVF)
+    assignSlot(Plan.RuntimeVF);
+
   ReversePostOrderTraversal<
       VPBlockRecursiveTraversalWrapper<const VPBlockBase *>>
       RPOT(VPBlockRecursiveTraversalWrapper<const VPBlockBase *>(
           Plan.getEntry()));
   for (const VPBasicBlock *VPBB :
        VPBlockUtils::blocksOnly<const VPBasicBlock>(RPOT))
     for (const VPRecipeBase &Recipe : *VPBB)
       for (VPValue *Def : Recipe.definedValues())
         assignSlot(Def);
 }