[VPlan] NFC: Change VFRange to take ElementCount

[sdesmalen]

Summary

This patch changes the type of Start, End in VFRange to be an ElementCount
instead of unsigned. This is done as preparation to make VPlans for
scalable vectors, but is otherwise NFC.

[dmgreen]

This looks simple enough to me. I guess it does mean we can't use the same vplan for representing scalable and non-scalable vectorization plans, but that's probably a sensible thing to do.

LGTM.

[fhahn]

LGTM

llvm/lib/Transforms/Vectorize/VPlan.h

75

Now that we have a constructor here, can we also assert that Start is a power-of-2?

[vkmr]

Other than the minor nit, LGTM.

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

7173

Nit: Add an isEmpty() method to VFRange. Removes the burden of knowing the condition for an empty range.

[sdesmalen]

• Added isEmpty method to VFRange.
• Added check that VFRange::Start is a power of 2.

[sdesmalen]

Thanks for the reviews and suggestions!

[vkmr]

LGTM.

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

if (!IsUnrollMetadata) {
     // Set operand 0 to refer to the loop id itself.
     NewLoopID->replaceOperandWith(0, NewLoopID);
     L->setLoopID(NewLoopID);
   }
 }
 
 bool LoopVectorizationPlanner::getDecisionAndClampRange(
     const std::function<bool(ElementCount)> &Predicate, VFRange &Range) {
-  assert(Range.End > Range.Start && "Trying to test an empty VF range.");
-  bool PredicateAtRangeStart = Predicate(ElementCount::getFixed(Range.Start));
-
-  for (unsigned TmpVF = Range.Start * 2; TmpVF < Range.End; TmpVF *= 2)
-    if (Predicate(ElementCount::getFixed(TmpVF)) != PredicateAtRangeStart) {
+  assert(ElementCount::isKnownGT(Range.End, Range.Start) &&
+         "Trying to test an empty VF range.");
+  bool PredicateAtRangeStart = Predicate(Range.Start);
+
+  for (ElementCount TmpVF = Range.Start * 2;
+       ElementCount::isKnownLT(TmpVF, Range.End); TmpVF *= 2)
+    if (Predicate(TmpVF) != PredicateAtRangeStart) {
       Range.End = TmpVF;
       break;
     }
 
   return PredicateAtRangeStart;
 }
 
 /// Build VPlans for the full range of feasible VF's = {\p MinVF, 2 * \p MinVF,
 /// 4 * \p MinVF, ..., \p MaxVF} by repeatedly building a VPlan for a sub-range
 /// of VF's starting at a given VF and extending it as much as possible. Each
 /// vectorization decision can potentially shorten this sub-range during
 /// buildVPlan().
 void LoopVectorizationPlanner::buildVPlans(unsigned MinVF, unsigned MaxVF) {
-  for (unsigned VF = MinVF; VF < MaxVF + 1;) {
-    VFRange SubRange = {VF, MaxVF + 1};
+  auto MaxVFPlusOne = ElementCount::getFixed(MaxVF).add(1);
+  for (ElementCount VF = ElementCount::getFixed(MinVF);
+       ElementCount::isKnownLT(VF, MaxVFPlusOne);) {
+    VFRange SubRange = {VF, MaxVFPlusOne};
     VPlans.push_back(buildVPlan(SubRange));
     VF = SubRange.End;
   }
 }
 
 VPValue *VPRecipeBuilder::createEdgeMask(BasicBlock *Src, BasicBlock *Dst,
                                          VPlanPtr &Plan) {
   assert(is_contained(predecessors(Dst), Src) && "Invalid edge");
void LoopVectorizationPlanner::buildVPlansWithVPRecipes(unsigned MinVF,
   auto &ConditionalAssumes = Legal->getConditionalAssumes();
   DeadInstructions.insert(ConditionalAssumes.begin(), ConditionalAssumes.end());
 
   DenseMap<Instruction *, Instruction *> &SinkAfter = Legal->getSinkAfter();
   // Dead instructions do not need sinking. Remove them from SinkAfter.
   for (Instruction *I : DeadInstructions)
     SinkAfter.erase(I);
 
-  for (unsigned VF = MinVF; VF < MaxVF + 1;) {
-    VFRange SubRange = {VF, MaxVF + 1};
+  auto MaxVFPlusOne = ElementCount::getFixed(MaxVF).add(1);
+  for (ElementCount VF = ElementCount::getFixed(MinVF);
+       ElementCount::isKnownLT(VF, MaxVFPlusOne);) {
+    VFRange SubRange = {VF, MaxVFPlusOne};
     VPlans.push_back(buildVPlanWithVPRecipes(SubRange, NeedDef,
                                              DeadInstructions, SinkAfter));
     VF = SubRange.End;
   }
 }
 
 VPlanPtr LoopVectorizationPlanner::buildVPlanWithVPRecipes(
     VFRange &Range, SmallPtrSetImpl<Value *> &NeedDef,
for (auto IG : InterleaveGroups) {
 
     for (unsigned i = 0; i < IG->getFactor(); ++i)
       if (Instruction *Member = IG->getMember(i)) {
         RecipeBuilder.getRecipe(Member)->eraseFromParent();
       }
   }
 
   // Adjust the recipes for any inloop reductions.
-  if (Range.Start > 1)
+  if (Range.Start.isVector())
     adjustRecipesForInLoopReductions(Plan, RecipeBuilder);
 
   // Finally, if tail is folded by masking, introduce selects between the phi
   // and the live-out instruction of each reduction, at the end of the latch.
   if (CM.foldTailByMasking() && !Legal->getReductionVars().empty()) {
     Builder.setInsertPoint(VPBB);
     auto *Cond = RecipeBuilder.createBlockInMask(OrigLoop->getHeader(), Plan);
     for (auto &Reduction : Legal->getReductionVars()) {
       if (CM.isInLoopReduction(Reduction.first))
         continue;
       VPValue *Phi = Plan->getVPValue(Reduction.first);
       VPValue *Red = Plan->getVPValue(Reduction.second.getLoopExitInstr());
       Builder.createNaryOp(Instruction::Select, {Cond, Red, Phi});
     }
   }
 
   std::string PlanName;
   raw_string_ostream RSO(PlanName);
-  ElementCount VF = ElementCount::getFixed(Range.Start);
+  ElementCount VF = Range.Start;
   Plan->addVF(VF);
   RSO << "Initial VPlan for VF={" << VF;
-  for (VF *= 2; VF.getKnownMinValue() < Range.End; VF *= 2) {
+  for (VF *= 2; ElementCount::isKnownLT(VF, Range.End); VF *= 2) {
     Plan->addVF(VF);
     RSO << "," << VF;
   }
   RSO << "},UF>=1";
   RSO.flush();
   Plan->setName(PlanName);
 
   return Plan;
VPlanPtr LoopVectorizationPlanner::buildVPlan(VFRange &Range) {
 
   // Create new empty VPlan
   auto Plan = std::make_unique<VPlan>();
 
   // Build hierarchical CFG
   VPlanHCFGBuilder HCFGBuilder(OrigLoop, LI, *Plan);
   HCFGBuilder.buildHierarchicalCFG();
 
-  for (unsigned VF = Range.Start; VF < Range.End; VF *= 2)
-    Plan->addVF(ElementCount::getFixed(VF));
+  for (ElementCount VF = Range.Start; ElementCount::isKnownLT(VF, Range.End);
+       VF *= 2)
+    Plan->addVF(VF);
 
   if (EnableVPlanPredication) {
     VPlanPredicator VPP(*Plan);
     VPP.predicate();
 
     // Avoid running transformation to recipes until masked code generation in
     // VPlan-native path is in place.
     return Plan;

llvm/lib/Transforms/Vectorize/VPlan.h

 class VPlan;
 class VPlanSlp;
 
 /// A range of powers-of-2 vectorization factors with fixed start and
 /// adjustable end. The range includes start and excludes end, e.g.,:
 /// [1, 9) = {1, 2, 4, 8}
 struct VFRange {
   // A power of 2.
-  const unsigned Start;
+  const ElementCount Start;
 
   // Need not be a power of 2. If End <= Start range is empty.
-  unsigned End;
+  ElementCount End;
+
+  VFRange(const ElementCount &Start, const ElementCount &End)
+      : Start(Start), End(End) {
+    assert(Start.isScalable() == End.isScalable() &&
+           "Both Start and End should have the same scalable flag");
+  }
 };
 
 using VPlanPtr = std::unique_ptr<VPlan>;
 
 /// In what follows, the term "input IR" refers to code that is fed into the
 /// vectorizer whereas the term "output IR" refers to code that is generated by
 /// the vectorizer.