[LV] Expose the primary IV as part of the plan, lit trip count (NFC).

[fhahn]

Summary

VPRecipeBuilder::createBlockInMask currently directly uses
Legal->getPrimaryInduction(). This tightly couples mask code-generation
with Legal. This patch introduces a new PrimaryIV VPValue to the VPlan
(similar to the BackedgeTakenCount), which can be used for masking,
instead of querying legal directly.

This uncouples code-generation from Legal and allows for changing the
primary IV before code-generation.

[fhahn]

Alternative approach in D7763 looks like the better way forward

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

 //===- LoopVectorize.cpp - A Loop Vectorizer ------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
void LoopVectorizationPlanner::executePlan(InnerLoopVectorizer &ILV,
   // 1. Create a new empty loop. Unlink the old loop and connect the new one.
   VPCallbackILV CallbackILV(ILV);
 
   VPTransformState State{BestVF, BestUF,      LI,
                          DT,     ILV.Builder, ILV.VectorLoopValueMap,
                          &ILV,   CallbackILV};
   State.CFG.PrevBB = ILV.createVectorizedLoopSkeleton();
   State.TripCount = ILV.getOrCreateTripCount(nullptr);
+  State.PrimaryIV = ILV.Legal->getPrimaryInduction();
 
   //===------------------------------------------------===//
   //
   // Notice: any optimization or new instruction that go
   // into the code below should also be implemented in
   // the cost-model.
   //
   //===------------------------------------------------===//
VPValue *VPRecipeBuilder::createBlockInMask(BasicBlock *BB, VPlanPtr &Plan) {
   VPValue *BlockMask = nullptr;
 
   if (OrigLoop->getHeader() == BB) {
     if (!CM.blockNeedsPredication(BB))
       return BlockMaskCache[BB] = BlockMask; // Loop incoming mask is all-one.
 
     // 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.
-    VPValue *IV = Plan->getVPValue(Legal->getPrimaryInduction());
+    VPValue *IV = Plan->getOrCreatePrimaryIV();
     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);

llvm/lib/Transforms/Vectorize/VPlan.h

 //===- VPlan.h - Represent A Vectorizer Plan --------------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
struct VPTransformState {
 
   /// Hold a reference to a mapping between VPValues in VPlan and original
   /// Values they correspond to.
   VPValue2ValueTy VPValue2Value;
 
   /// Hold the trip count of the scalar loop.
   Value *TripCount = nullptr;
 
+  Value *PrimaryIV = nullptr;
+
   /// Hold a pointer to InnerLoopVectorizer to reuse its IR generation methods.
   InnerLoopVectorizer *ILV;
 
   VPCallback &Callback;
 };
 
 /// VPBlockBase is the building block of the Hierarchical Control-Flow Graph.
 /// A VPBlockBase can be either a VPBasicBlock or a VPRegionBlock.
private:
   // underlying IR that will be used to implement its structural comparison
   // (operators '==' and '<').
   SmallPtrSet<VPValue *, 16> VPExternalDefs;
 
   /// Represents the backedge taken count of the original loop, for folding
   /// the tail.
   VPValue *BackedgeTakenCount = nullptr;
 
+  /// Represents the primary IV of the original loop.
+  VPValue *PrimaryIV = nullptr;
+
   /// Holds a mapping between Values and their corresponding VPValue inside
   /// VPlan.
   Value2VPValueTy Value2VPValue;
 
   /// Holds the VPLoopInfo analysis for this VPlan.
   VPLoopInfo VPLInfo;
 
   /// Holds the condition bit values built during VPInstruction to VPRecipe transformation.
   SmallVector<VPValue *, 4> VPCBVs;
 
 public:
   VPlan(VPBlockBase *Entry = nullptr) : Entry(Entry) {
     if (Entry)
       Entry->setPlan(this);
   }
 
   ~VPlan() {
     if (Entry)
       VPBlockBase::deleteCFG(Entry);
     for (auto &MapEntry : Value2VPValue)
-      if (MapEntry.second != BackedgeTakenCount)
+      if (MapEntry.second != BackedgeTakenCount && MapEntry.second != PrimaryIV)
         delete MapEntry.second;
     if (BackedgeTakenCount)
       delete BackedgeTakenCount; // Delete once, if in Value2VPValue or not.
+    if (PrimaryIV)
+      delete PrimaryIV; // Delete once, if in Value2VPValue or not.
     for (VPValue *Def : VPExternalDefs)
       delete Def;
     for (VPValue *CBV : VPCBVs)
       delete CBV;
   }
 
   /// Generate the IR code for this VPlan.
   void execute(struct VPTransformState *State);
public:
 
   /// The backedge taken count of the original loop.
   VPValue *getOrCreateBackedgeTakenCount() {
     if (!BackedgeTakenCount)
       BackedgeTakenCount = new VPValue();
     return BackedgeTakenCount;
   }
 
+  VPValue *getOrCreatePrimaryIV() {
+    if (!PrimaryIV)
+      PrimaryIV = new VPValue();
+    return PrimaryIV;
+  }
+
   void addVF(unsigned VF) { VFs.insert(VF); }
 
   bool hasVF(unsigned 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

 //===- VPlan.cpp - Vectorizer Plan ----------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 ///
void VPlan::execute(VPTransformState *State) {
   // -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");
     Value2VPValue[TCMO] = BackedgeTakenCount;
   }
+  if (PrimaryIV && PrimaryIV->getNumUsers())
+    Value2VPValue[State->PrimaryIV] = PrimaryIV;
 
   // 0. Set the reverse mapping from VPValues to Values for code generation.
   for (auto &Entry : Value2VPValue)
     State->VPValue2Value[Entry.second] = Entry.first;
 
   BasicBlock *VectorPreHeaderBB = State->CFG.PrevBB;
   BasicBlock *VectorHeaderBB = VectorPreHeaderBB->getSingleSuccessor();
   assert(VectorHeaderBB && "Loop preheader does not have a single successor.");