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 @@ -1720,6 +1720,36 @@ } }; +/// Recipe to expand a SCEV expression. +/// TODO: Currently the recipe always expands the expression in the loop +/// pre-header, but the recipe is currently placed in the header; place it in +/// the pre-header once the latter is modeled in VPlan as a VPBasicBlock. +class VPExpandSCEVRecipe : public VPRecipeBase, public VPValue { + const SCEV *Expr; + ScalarEvolution &SE; + +public: + VPExpandSCEVRecipe(const SCEV *Expr, ScalarEvolution &SE) + : VPRecipeBase(VPExpandSCEVSC, {}), VPValue(nullptr, this), Expr(Expr), + SE(SE) {} + + ~VPExpandSCEVRecipe() override = default; + + /// Method to support type inquiry through isa, cast, and dyn_cast. + static inline bool classof(const VPDef *D) { + return D->getVPDefID() == VPExpandSCEVSC; + } + + /// Generate a canonical vector induction variable of the vector loop, with + void execute(VPTransformState &State) override; + +#if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP) + /// Print the recipe. + void print(raw_ostream &O, const Twine &Indent, + VPSlotTracker &SlotTracker) const override; +#endif +}; + /// Canonical scalar induction phi of the vector loop. Starting at the specified /// start value (either 0 or the resume value when vectorizing the epilogue /// loop). VPWidenCanonicalIVRecipe represents the vector version of the diff --git a/llvm/lib/Transforms/Vectorize/VPlan.cpp b/llvm/lib/Transforms/Vectorize/VPlan.cpp --- a/llvm/lib/Transforms/Vectorize/VPlan.cpp +++ b/llvm/lib/Transforms/Vectorize/VPlan.cpp @@ -41,6 +41,7 @@ #include "llvm/Support/GraphWriter.h" #include "llvm/Support/raw_ostream.h" #include "llvm/Transforms/Utils/BasicBlockUtils.h" +#include "llvm/Transforms/Utils/ScalarEvolutionExpander.h" #include #include #include @@ -1399,6 +1400,26 @@ } #endif +void VPExpandSCEVRecipe::execute(VPTransformState &State) { + assert(!State.Instance && "cannot be used in per-lane"); + const DataLayout &DL = + State.CFG.VectorPreHeader->getModule()->getDataLayout(); + SCEVExpander Exp(SE, DL, "induction"); + auto Res = Exp.expandCodeFor(Expr, Expr->getType(), + State.CFG.VectorPreHeader->getTerminator()); + + State.set(this, Res, 0); +} + +#if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP) +void VPExpandSCEVRecipe::print(raw_ostream &O, const Twine &Indent, + VPSlotTracker &SlotTracker) const { + O << Indent << "EMIT "; + getVPSingleValue()->printAsOperand(O, SlotTracker); + O << " = EXPAND SCEV " << *Expr; +} +#endif + void VPWidenCanonicalIVRecipe::execute(VPTransformState &State) { Value *CanonicalIV = State.get(getOperand(0), 0); Type *STy = CanonicalIV->getType(); diff --git a/llvm/lib/Transforms/Vectorize/VPlanValue.h b/llvm/lib/Transforms/Vectorize/VPlanValue.h --- a/llvm/lib/Transforms/Vectorize/VPlanValue.h +++ b/llvm/lib/Transforms/Vectorize/VPlanValue.h @@ -327,6 +327,7 @@ /// type identification. using VPRecipeTy = enum { VPBranchOnMaskSC, + VPExpandSCEVSC, VPInstructionSC, VPInterleaveSC, VPReductionSC,