[VPlan] Add InductionDescriptor to VPWidenIntOrFpInduction. (NFC)

[fhahn]

Summary

This allows easier access to the induction descriptor from VPlan,
without needing to go through Legal. VPReductionPHIRecipe already
contains a RecurrenceDescriptor in a similar fashion.

[Ayal]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

8591

This change from lookup() to find() is done in order to obtain a reference to an existing InductionDescriptor to be recorded in VPWidenIntOrFpInduction recipe, rather than a copy thereof. Right?

Can alternatively do

if (!Legal->isInductionPhi(Phi))
  return nullptr;
InductionDescriptor &II = Legal->getInductionVars().find(Phi)->second;

8623

This is a similar setting of II as suggested above, but where checking isInductionPhi() is redundant.

llvm/lib/Transforms/Vectorize/VPlan.h

65

Lex order

1068

IndDesc >> Descriptor?
Or, to better match VPReductionPHIRecipe, getDescriptor() >> getInductionDescriptor()?

llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp

49

A similar case, albeit w/o direct access to Legal.
Should Legal support all above cases by providing something like:

Optional<InductionDescriptor *> getInductionDescriptor(PHINode *Phi) {
  if (!isInductionPhi(Phi))
    return None;
  return &getInductionVars().find(Phi)->second;
}

or even

Optional<InductionDescriptor *> getIntOrFpInductionDescriptor(PHINode *Phi) {
  if (!isInductionPhi(Phi))
    return None;
  auto &ID = getInductionVars().find(Phi)->second;
  if (ID.getKind() == InductionDescriptor::IK_IntInduction ||
      ID.getKind() == InductionDescriptor::IK_FpInduction)
    return &ID;
  return None;
}

?

[fhahn]

Adjust naming, add and use getIntOrFpInductionDescriptor helper as suggested.

[fhahn]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

8591

I went ahead and added a getIntOrFpInductionDescriptor helper as suggested below.

llvm/lib/Transforms/Vectorize/VPlan.h

65

Moved to correct position.

1068

updated to getInductionDescriptor.

llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp

49

I went ahead and added a getIntOrFpInductionDescriptor helper. It is used in LV and here. VPInstructionsToVPRecipes has been updated to take a function_ref to get use it. That should reduce duplication to a minimum

[Ayal]

Thanks! Ship it, just a couple of typos to fix.

llvm/include/llvm/Transforms/Vectorize/LoopVectorizationLegality.h

313

[f]or

llvm/lib/Transforms/Vectorize/VPlan.h

1067

to >> the

llvm/lib/Transforms/Vectorize/VPlanTransforms.h

26

Lex order.

[fhahn]

llvm/include/llvm/Transforms/Vectorize/LoopVectorizationLegality.h

313

Thanks, should be fixed in the committed version.

llvm/lib/Transforms/Vectorize/VPlan.h

1067

Thanks, should be fixed in the committed version.

llvm/lib/Transforms/Vectorize/VPlanTransforms.h

26

Thanks, should be fixed in the committed version.

llvm/include/llvm/Transforms/Vectorize/LoopVectorizationLegality.h

public:
   MapVector<Instruction *, Instruction *> &getSinkAfter() { return SinkAfter; }
 
   /// Returns the widest induction type.
   Type *getWidestInductionType() { return WidestIndTy; }
 
   /// Returns True if V is a Phi node of an induction variable in this loop.
   bool isInductionPhi(const Value *V) const;
 
+  /// Returns a pointer to the induction descriptor, if \p Phi is an integer or
+  /// floating point induction.
+  const InductionDescriptor *getIntOrFpInductionDescriptor(PHINode *Phi) const;
+
   /// Returns True if V is a cast that is part of an induction def-use chain,
   /// and had been proven to be redundant under a runtime guard (in other
   /// words, the cast has the same SCEV expression as the induction phi).
   bool isCastedInductionVariable(const Value *V) const;
 
   /// Returns True if V can be considered as an induction variable in this
   /// loop. V can be the induction phi, or some redundant cast in the def-use
   /// chain of the inducion phi.

llvm/lib/Transforms/Vectorize/LoopVectorizationLegality.cpp

bool LoopVectorizationLegality::isInductionPhi(const Value *V) const {
   Value *In0 = const_cast<Value *>(V);
   PHINode *PN = dyn_cast_or_null<PHINode>(In0);
   if (!PN)
     return false;
 
   return Inductions.count(PN);
 }
 
+const InductionDescriptor *
+LoopVectorizationLegality::getIntOrFpInductionDescriptor(PHINode *Phi) const {
+  if (!isInductionPhi(Phi))
+    return nullptr;
+  auto &ID = getInductionVars().find(Phi)->second;
+  if (ID.getKind() == InductionDescriptor::IK_IntInduction ||
+      ID.getKind() == InductionDescriptor::IK_FpInduction)
+    return &ID;
+  return nullptr;
+}
+
 bool LoopVectorizationLegality::isCastedInductionVariable(
     const Value *V) const {
   auto *Inst = dyn_cast<Instruction>(V);
   return (Inst && InductionCastsToIgnore.count(Inst));
 }
 
 bool LoopVectorizationLegality::isInductionVariable(const Value *V) const {
   return isInductionPhi(V) || isCastedInductionVariable(V);

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

public:
   /// Instr's operands.
   void scalarizeInstruction(Instruction *Instr, VPReplicateRecipe *RepRecipe,
                             const VPIteration &Instance, bool IfPredicateInstr,
                             VPTransformState &State);
 
   /// Widen an integer or floating-point induction variable \p IV. If \p Trunc
   /// is provided, the integer induction variable will first be truncated to
   /// the corresponding type.
-  void widenIntOrFpInduction(PHINode *IV, Value *Start, TruncInst *Trunc,
-                             VPValue *Def, VPValue *CastDef,
-                             VPTransformState &State);
+  void widenIntOrFpInduction(PHINode *IV, const InductionDescriptor &ID,
+                             Value *Start, TruncInst *Trunc, VPValue *Def,
+                             VPValue *CastDef, VPTransformState &State);
 
   /// Construct the vector value of a scalarized value \p V one lane at a time.
   void packScalarIntoVectorValue(VPValue *Def, const VPIteration &Instance,
                                  VPTransformState &State);
 
   /// Try to vectorize interleaved access group \p Group with the base address
   /// given in \p Addr, optionally masking the vector operations if \p
   /// BlockInMask is non-null. Use \p State to translate given VPValues to IR
void InnerLoopVectorizer::recordVectorLoopValueForInductionCast(
   // The rest of the Casts (if exist) have no uses outside the
   // induction update chain itself.
   if (Lane < UINT_MAX)
     State.set(CastDef, VectorLoopVal, VPIteration(Part, Lane));
   else
     State.set(CastDef, VectorLoopVal, Part);
 }
 
-void InnerLoopVectorizer::widenIntOrFpInduction(PHINode *IV, Value *Start,
-                                                TruncInst *Trunc, VPValue *Def,
-                                                VPValue *CastDef,
+void InnerLoopVectorizer::widenIntOrFpInduction(PHINode *IV,
+                                                const InductionDescriptor &ID,
+                                                Value *Start, TruncInst *Trunc,
+                                                VPValue *Def, VPValue *CastDef,
                                                 VPTransformState &State) {
   assert((IV->getType()->isIntegerTy() || IV != OldInduction) &&
          "Primary induction variable must have an integer type");
-
-  auto II = Legal->getInductionVars().find(IV);
-  assert(II != Legal->getInductionVars().end() && "IV is not an induction");
-
-  auto ID = II->second;
   assert(IV->getType() == ID.getStartValue()->getType() && "Types must match");
 
   // The value from the original loop to which we are mapping the new induction
   // variable.
   Instruction *EntryVal = Trunc ? cast<Instruction>(Trunc) : IV;
 
   auto &DL = OrigLoop->getHeader()->getModule()->getDataLayout();
 
return new VPWidenMemoryInstructionRecipe(*Store, Operands[1], Operands[0],
                                             Mask, Consecutive, Reverse);
 }
 
 VPWidenIntOrFpInductionRecipe *
 VPRecipeBuilder::tryToOptimizeInductionPHI(PHINode *Phi,
                                            ArrayRef<VPValue *> Operands) const {
   // Check if this is an integer or fp induction. If so, build the recipe that
   // produces its scalar and vector values.
-  InductionDescriptor II = Legal->getInductionVars().lookup(Phi);
-  if (II.getKind() == InductionDescriptor::IK_IntInduction ||
-      II.getKind() == InductionDescriptor::IK_FpInduction) {
-    assert(II.getStartValue() ==
+  if (auto *II = Legal->getIntOrFpInductionDescriptor(Phi)) {
+    assert(II->getStartValue() ==
            Phi->getIncomingValueForBlock(OrigLoop->getLoopPreheader()));
-    const SmallVectorImpl<Instruction *> &Casts = II.getCastInsts();
+    const SmallVectorImpl<Instruction *> &Casts = II->getCastInsts();
     return new VPWidenIntOrFpInductionRecipe(
-        Phi, Operands[0], Casts.empty() ? nullptr : Casts.front());
+        Phi, Operands[0], *II, Casts.empty() ? nullptr : Casts.front());
   }
 
   return nullptr;
 }
 
 VPWidenIntOrFpInductionRecipe *VPRecipeBuilder::tryToOptimizeInductionTruncate(
     TruncInst *I, ArrayRef<VPValue *> Operands, VFRange &Range,
     VPlan &Plan) const {
auto isOptimizableIVTruncate =
     return [=](ElementCount VF) -> bool {
       return CM.isOptimizableIVTruncate(K, VF);
     };
   };
 
   if (LoopVectorizationPlanner::getDecisionAndClampRange(
           isOptimizableIVTruncate(I), Range)) {
 
-    InductionDescriptor II =
-        Legal->getInductionVars().lookup(cast<PHINode>(I->getOperand(0)));
+    auto *Phi = cast<PHINode>(I->getOperand(0));
+    const InductionDescriptor &II = *Legal->getIntOrFpInductionDescriptor(Phi);
     VPValue *Start = Plan.getOrAddVPValue(II.getStartValue());
-    return new VPWidenIntOrFpInductionRecipe(cast<PHINode>(I->getOperand(0)),
-                                             Start, I);
+    return new VPWidenIntOrFpInductionRecipe(Phi, Start, II, I);
   }
   return nullptr;
 }
 
 VPRecipeOrVPValueTy VPRecipeBuilder::tryToBlend(PHINode *Phi,
                                                 ArrayRef<VPValue *> Operands,
                                                 VPlanPtr &Plan) {
   // If all incoming values are equal, the incoming VPValue can be used directly
if (EnableVPlanPredication) {
     VPP.predicate();
 
     // Avoid running transformation to recipes until masked code generation in
     // VPlan-native path is in place.
     return Plan;
   }
 
   SmallPtrSet<Instruction *, 1> DeadInstructions;
-  VPlanTransforms::VPInstructionsToVPRecipes(OrigLoop, Plan,
-                                             Legal->getInductionVars(),
-                                             DeadInstructions, *PSE.getSE());
+  VPlanTransforms::VPInstructionsToVPRecipes(
+      OrigLoop, Plan,
+      [this](PHINode *P) { return Legal->getIntOrFpInductionDescriptor(P); },
   return Plan;
 }
 
 // Adjust the recipes for reductions. For in-loop reductions the chain of
 // instructions leading from the loop exit instr to the phi need to be converted
 // to reductions, with one operand being vector and the other being the scalar
 // reduction chain. For other reductions, a select is introduced between the phi
 // and live-out recipes when folding the tail.
for (unsigned Part = 0; Part < State.UF; ++Part) {
       State.set(this, NewGEP, Part);
       State.ILV->addMetadata(NewGEP, GEP);
     }
   }
 }
 
 void VPWidenIntOrFpInductionRecipe::execute(VPTransformState &State) {
   assert(!State.Instance && "Int or FP induction being replicated.");
-  State.ILV->widenIntOrFpInduction(IV, getStartValue()->getLiveInIRValue(),
-                                   getTruncInst(), getVPValue(0),
-                                   getCastValue(), State);
+  State.ILV->widenIntOrFpInduction(
+      IV, getInductionDescriptor(), getStartValue()->getLiveInIRValue(),
+      getTruncInst(), getVPValue(0), getCastValue(), State);
 }
 
 void VPWidenPHIRecipe::execute(VPTransformState &State) {
   State.ILV->widenPHIInstruction(cast<PHINode>(getUnderlyingValue()), this,
                                  State);
 }
 
 void VPBlendRecipe::execute(VPTransformState &State) {

llvm/lib/Transforms/Vectorize/VPlan.h

 #include <cstddef>
 #include <map>
 #include <string>
 
 namespace llvm {
 
 class BasicBlock;
 class DominatorTree;
+class InductionDescriptor;
 class InnerLoopVectorizer;
 class LoopInfo;
 class raw_ostream;
 class RecurrenceDescriptor;
 class Value;
 class VPBasicBlock;
 class VPRegionBlock;
 class VPlan;
 class VPReplicateRecipe;
 class VPlanSlp;
 
 /// Returns a calculation for the total number of elements for a given \p VF.
 /// For fixed width vectors this value is a constant, whereas for scalable
 /// vectors it is an expression determined at runtime.
 Value *getRuntimeVF(IRBuilder<> &B, Type *Ty, ElementCount VF);
 
 /// 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}
void print(raw_ostream &O, const Twine &Indent,
              VPSlotTracker &SlotTracker) const override;
 #endif
 };
 
 /// A recipe for handling phi nodes of integer and floating-point inductions,
 /// producing their vector and scalar values.
 class VPWidenIntOrFpInductionRecipe : public VPRecipeBase {
   PHINode *IV;
+  const InductionDescriptor &IndDesc;
 
 public:
   VPWidenIntOrFpInductionRecipe(PHINode *IV, VPValue *Start,
+                                const InductionDescriptor &IndDesc,
                                 Instruction *Cast = nullptr)
-      : VPRecipeBase(VPWidenIntOrFpInductionSC, {Start}), IV(IV) {
+      : VPRecipeBase(VPWidenIntOrFpInductionSC, {Start}), IV(IV),
+        IndDesc(IndDesc) {
     new VPValue(IV, this);
 
     if (Cast)
       new VPValue(Cast, this);
   }
 
-  VPWidenIntOrFpInductionRecipe(PHINode *IV, VPValue *Start, TruncInst *Trunc)
-      : VPRecipeBase(VPWidenIntOrFpInductionSC, {Start}), IV(IV) {
+  VPWidenIntOrFpInductionRecipe(PHINode *IV, VPValue *Start,
+                                const InductionDescriptor &IndDesc,
+                                TruncInst *Trunc)
+      : VPRecipeBase(VPWidenIntOrFpInductionSC, {Start}), IV(IV),
+        IndDesc(IndDesc) {
     new VPValue(Trunc, this);
   }
 
   ~VPWidenIntOrFpInductionRecipe() override = default;
 
   /// Method to support type inquiry through isa, cast, and dyn_cast.
   static inline bool classof(const VPDef *D) {
     return D->getVPDefID() == VPRecipeBase::VPWidenIntOrFpInductionSC;
#endif
   /// Returns the first defined value as TruncInst, if it is one or nullptr
   /// otherwise.
   TruncInst *getTruncInst() {
     return dyn_cast_or_null<TruncInst>(getVPValue(0)->getUnderlyingValue());
   }
   const TruncInst *getTruncInst() const {
     return dyn_cast_or_null<TruncInst>(getVPValue(0)->getUnderlyingValue());
   }
+
+  /// Returns the induction descriptor for the recipe.
+  const InductionDescriptor &getInductionDescriptor() const { return IndDesc; }
 };
 
 /// A recipe for handling first order recurrences and pointer inductions. For
 /// first-order recurrences, the start value is the first operand of the recipe
 /// and the incoming value from the backedge is the second operand. It also
 /// serves as base class for VPReductionPHIRecipe. In the VPlan native path, all
 /// incoming VPValues & VPBasicBlock pairs are managed in the recipe directly.
 class VPWidenPHIRecipe : public VPRecipeBase, public VPValue {

llvm/lib/Transforms/Vectorize/VPlanTransforms.h

 //===- VPlanTransforms.h - Utility VPlan to VPlan transforms --------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 ///
 /// \file
 /// This file provides utility VPlan to VPlan transformations.
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_TRANSFORMS_VECTORIZE_VPLANTRANSFORMS_H
 #define LLVM_TRANSFORMS_VECTORIZE_VPLANTRANSFORMS_H
 
 #include "VPlan.h"
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/Transforms/Vectorize/LoopVectorizationLegality.h"
 
 namespace llvm {
 
+class InductionDescriptor;
 class Instruction;
+class PHINode;
 class ScalarEvolution;
 
 struct VPlanTransforms {
   /// Replaces the VPInstructions in \p Plan with corresponding
   /// widen recipes.
-  static void VPInstructionsToVPRecipes(
-      Loop *OrigLoop, VPlanPtr &Plan,
-      const LoopVectorizationLegality::InductionList &Inductions,
-      SmallPtrSetImpl<Instruction *> &DeadInstructions, ScalarEvolution &SE);
+  static void
+  VPInstructionsToVPRecipes(Loop *OrigLoop, VPlanPtr &Plan,
+                            function_ref<const InductionDescriptor *(PHINode *)>
+                                GetIntOrFpInductionDescriptor,
+                            SmallPtrSetImpl<Instruction *> &DeadInstructions,
+                            ScalarEvolution &SE);
 
   static bool sinkScalarOperands(VPlan &Plan);
 
   static bool mergeReplicateRegions(VPlan &Plan);
 };
 
 } // namespace llvm
 
 #endif // LLVM_TRANSFORMS_VECTORIZE_VPLANTRANSFORMS_H

llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp

 
 #include "VPlanTransforms.h"
 #include "llvm/ADT/PostOrderIterator.h"
 
 using namespace llvm;
 
 void VPlanTransforms::VPInstructionsToVPRecipes(
     Loop *OrigLoop, VPlanPtr &Plan,
-    const LoopVectorizationLegality::InductionList &Inductions,
+    function_ref<const InductionDescriptor *(PHINode *)>
+        GetIntOrFpInductionDescriptor,
     SmallPtrSetImpl<Instruction *> &DeadInstructions, ScalarEvolution &SE) {
 
   auto *TopRegion = cast<VPRegionBlock>(Plan->getEntry());
   ReversePostOrderTraversal<VPBlockBase *> RPOT(TopRegion->getEntry());
 
   for (VPBlockBase *Base : RPOT) {
     // Do not widen instructions in pre-header and exit blocks.
     if (Base->getNumPredecessors() == 0 || Base->getNumSuccessors() == 0)
for (VPRecipeBase &Ingredient : llvm::make_early_inc_range(*VPBB)) {
         VPV->replaceAllUsesWith(&DummyValue);
         Ingredient.eraseFromParent();
         continue;
       }
 
       VPRecipeBase *NewRecipe = nullptr;
       if (auto *VPPhi = dyn_cast<VPWidenPHIRecipe>(&Ingredient)) {
         auto *Phi = cast<PHINode>(VPPhi->getUnderlyingValue());
-        InductionDescriptor II = Inductions.lookup(Phi);
-        if (II.getKind() == InductionDescriptor::IK_IntInduction ||
-            II.getKind() == InductionDescriptor::IK_FpInduction) {
-          VPValue *Start = Plan->getOrAddVPValue(II.getStartValue());
-          NewRecipe = new VPWidenIntOrFpInductionRecipe(Phi, Start);
+        if (const auto *II = GetIntOrFpInductionDescriptor(Phi)) {
+          VPValue *Start = Plan->getOrAddVPValue(II->getStartValue());
+          NewRecipe = new VPWidenIntOrFpInductionRecipe(Phi, Start, *II);
         } else {
           Plan->addVPValue(Phi, VPPhi);
           continue;
         }
       } else {
         assert(isa<VPInstruction>(&Ingredient) &&
                "only VPInstructions expected here");
         assert(!isa<PHINode>(Inst) && "phis should be handled above");

llvm/unittests/Transforms/Vectorize/VPlanHCFGTest.cpp

N3 [label =
       "No successors\l"
     ]
   }
 }
 )";
   EXPECT_EQ(ExpectedStr, FullDump);
 #endif
 
-  LoopVectorizationLegality::InductionList Inductions;
   SmallPtrSet<Instruction *, 1> DeadInstructions;
-  VPlanTransforms::VPInstructionsToVPRecipes(LI->getLoopFor(LoopHeader), Plan,
-                                             Inductions, DeadInstructions, *SE);
+  VPlanTransforms::VPInstructionsToVPRecipes(
+      LI->getLoopFor(LoopHeader), Plan, [](PHINode *P) { return nullptr; },
+      DeadInstructions, *SE);
 }
 
 TEST_F(VPlanHCFGTest, testVPInstructionToVPRecipesInner) {
   const char *ModuleString =
       "define void @f(i32* %A, i64 %N) {\n"
       "entry:\n"
       "  br label %for.body\n"
       "for.body:\n"
const char *ModuleString =
       "}\n";
 
   Module &M = parseModule(ModuleString);
 
   Function *F = M.getFunction("f");
   BasicBlock *LoopHeader = F->getEntryBlock().getSingleSuccessor();
   auto Plan = buildHCFG(LoopHeader);
 
-  LoopVectorizationLegality::InductionList Inductions;
   SmallPtrSet<Instruction *, 1> DeadInstructions;
-  VPlanTransforms::VPInstructionsToVPRecipes(LI->getLoopFor(LoopHeader), Plan,
-                                             Inductions, DeadInstructions, *SE);
+  VPlanTransforms::VPInstructionsToVPRecipes(
+      LI->getLoopFor(LoopHeader), Plan, [](PHINode *P) { return nullptr; },
+      DeadInstructions, *SE);
 
   VPBlockBase *Entry = Plan->getEntry()->getEntryBasicBlock();
   EXPECT_NE(nullptr, Entry->getSingleSuccessor());
   EXPECT_EQ(0u, Entry->getNumPredecessors());
   EXPECT_EQ(1u, Entry->getNumSuccessors());
 
   VPBasicBlock *VecBB = Entry->getSingleSuccessor()->getEntryBasicBlock();
   EXPECT_EQ(7u, VecBB->size());