diff --git a/llvm/include/llvm/Transforms/Vectorize/LoopVectorizationLegality.h b/llvm/include/llvm/Transforms/Vectorize/LoopVectorizationLegality.h --- a/llvm/include/llvm/Transforms/Vectorize/LoopVectorizationLegality.h +++ b/llvm/include/llvm/Transforms/Vectorize/LoopVectorizationLegality.h @@ -310,6 +310,10 @@ /// 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). diff --git a/llvm/lib/Transforms/Vectorize/LoopVectorizationLegality.cpp b/llvm/lib/Transforms/Vectorize/LoopVectorizationLegality.cpp --- a/llvm/lib/Transforms/Vectorize/LoopVectorizationLegality.cpp +++ b/llvm/lib/Transforms/Vectorize/LoopVectorizationLegality.cpp @@ -938,6 +938,17 @@ 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(V); diff --git a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp --- a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp +++ b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp @@ -506,9 +506,9 @@ /// 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, @@ -2487,17 +2487,13 @@ 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 @@ -8592,14 +8588,12 @@ ArrayRef 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 &Casts = II.getCastInsts(); + const SmallVectorImpl &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; @@ -8625,11 +8619,10 @@ if (LoopVectorizationPlanner::getDecisionAndClampRange( isOptimizableIVTruncate(I), Range)) { - InductionDescriptor II = - Legal->getInductionVars().lookup(cast(I->getOperand(0))); + auto *Phi = cast(I->getOperand(0)); + const InductionDescriptor &II = *Legal->getIntOrFpInductionDescriptor(Phi); VPValue *Start = Plan.getOrAddVPValue(II.getStartValue()); - return new VPWidenIntOrFpInductionRecipe(cast(I->getOperand(0)), - Start, I); + return new VPWidenIntOrFpInductionRecipe(Phi, Start, II, I); } return nullptr; } @@ -9359,9 +9352,10 @@ } SmallPtrSet DeadInstructions; - VPlanTransforms::VPInstructionsToVPRecipes(OrigLoop, Plan, - Legal->getInductionVars(), - DeadInstructions, *PSE.getSE()); + VPlanTransforms::VPInstructionsToVPRecipes( + OrigLoop, Plan, + [this](PHINode *P) { return Legal->getIntOrFpInductionDescriptor(P); }, + DeadInstructions, *PSE.getSE()); return Plan; } @@ -9719,9 +9713,9 @@ 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) { diff --git a/llvm/lib/Transforms/Vectorize/VPlan.h b/llvm/lib/Transforms/Vectorize/VPlan.h --- a/llvm/lib/Transforms/Vectorize/VPlan.h +++ b/llvm/lib/Transforms/Vectorize/VPlan.h @@ -51,6 +51,7 @@ class BasicBlock; class DominatorTree; +class InductionDescriptor; class InnerLoopVectorizer; class LoopInfo; class raw_ostream; @@ -1005,19 +1006,25 @@ /// 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); } @@ -1056,6 +1063,9 @@ const TruncInst *getTruncInst() const { return dyn_cast_or_null(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 diff --git a/llvm/lib/Transforms/Vectorize/VPlanTransforms.h b/llvm/lib/Transforms/Vectorize/VPlanTransforms.h --- a/llvm/lib/Transforms/Vectorize/VPlanTransforms.h +++ b/llvm/lib/Transforms/Vectorize/VPlanTransforms.h @@ -14,20 +14,25 @@ #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 &DeadInstructions, ScalarEvolution &SE); + static void + VPInstructionsToVPRecipes(Loop *OrigLoop, VPlanPtr &Plan, + function_ref + GetIntOrFpInductionDescriptor, + SmallPtrSetImpl &DeadInstructions, + ScalarEvolution &SE); static bool sinkScalarOperands(VPlan &Plan); diff --git a/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp b/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp --- a/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp +++ b/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp @@ -18,7 +18,8 @@ void VPlanTransforms::VPInstructionsToVPRecipes( Loop *OrigLoop, VPlanPtr &Plan, - const LoopVectorizationLegality::InductionList &Inductions, + function_ref + GetIntOrFpInductionDescriptor, SmallPtrSetImpl &DeadInstructions, ScalarEvolution &SE) { auto *TopRegion = cast(Plan->getEntry()); @@ -44,11 +45,9 @@ VPRecipeBase *NewRecipe = nullptr; if (auto *VPPhi = dyn_cast(&Ingredient)) { auto *Phi = cast(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; diff --git a/llvm/unittests/Transforms/Vectorize/VPlanHCFGTest.cpp b/llvm/unittests/Transforms/Vectorize/VPlanHCFGTest.cpp --- a/llvm/unittests/Transforms/Vectorize/VPlanHCFGTest.cpp +++ b/llvm/unittests/Transforms/Vectorize/VPlanHCFGTest.cpp @@ -134,10 +134,10 @@ EXPECT_EQ(ExpectedStr, FullDump); #endif - LoopVectorizationLegality::InductionList Inductions; SmallPtrSet 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) { @@ -164,10 +164,10 @@ BasicBlock *LoopHeader = F->getEntryBlock().getSingleSuccessor(); auto Plan = buildHCFG(LoopHeader); - LoopVectorizationLegality::InductionList Inductions; SmallPtrSet 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());