diff --git a/llvm/include/llvm/Analysis/IVDescriptors.h b/llvm/include/llvm/Analysis/IVDescriptors.h index 654a5f10cea96..0c83814881161 100644 --- a/llvm/include/llvm/Analysis/IVDescriptors.h +++ b/llvm/include/llvm/Analysis/IVDescriptors.h @@ -28,6 +28,7 @@ class Loop; class PredicatedScalarEvolution; class ScalarEvolution; class SCEV; +class SCEVAddRecExpr; class StoreInst; /// These are the kinds of recurrences that we support. @@ -477,6 +478,44 @@ class InductionDescriptor { SmallVector RedundantCasts; }; +/// A struct for saving information about monotonic variables. +/// Monotonic variable can be considered as a "conditional" induction variable: +/// its update happens only on loop iterations for which a certain predicate is +/// satisfied. In this implementation the predicate is represented as an edge in +/// loop CFG: variable is updated if this edge is executed on current loop +/// iteration. +class MonotonicDescriptor { +public: + using Edge = std::pair; + + MonotonicDescriptor() = default; + + const SmallPtrSetImpl &getChain() const { return Chain; } + Instruction *getStepInst() const { return StepInst; } + Edge getPredicateEdge() const { return PredEdge; } + const SCEVAddRecExpr *getExpr() const { return Expr; } + + /// Returns true if \p PN is a monotonic variable in the loop \p L. If \p PN + /// is monotonic, the monotonic descriptor \p D will contain the data + /// describing this variable. + static bool isMonotonicPHI(PHINode *PN, const Loop *L, + MonotonicDescriptor &Desc, ScalarEvolution &SE); + + /// Returns true if \p Val is a monotonic variable in the loop \p L (in this + /// case, the value should transitively contain monotonic phi as part of its + /// calculation). + static bool isMonotonicVal(Value *Val, const Loop *L, + MonotonicDescriptor &Desc, ScalarEvolution &SE); + +private: + SmallPtrSet Chain; + Instruction *StepInst; + Edge PredEdge; + const SCEVAddRecExpr *Expr; + + bool setSCEV(const SCEV *NewExpr); +}; + } // end namespace llvm #endif // LLVM_ANALYSIS_IVDESCRIPTORS_H diff --git a/llvm/lib/Analysis/IVDescriptors.cpp b/llvm/lib/Analysis/IVDescriptors.cpp index 9f8ac6e8e2e0b..4048e32e07c2b 100644 --- a/llvm/lib/Analysis/IVDescriptors.cpp +++ b/llvm/lib/Analysis/IVDescriptors.cpp @@ -1661,3 +1661,124 @@ bool InductionDescriptor::isInductionPHI( D = InductionDescriptor(StartValue, IK_PtrInduction, Step); return true; } + +bool MonotonicDescriptor::setSCEV(const SCEV *NewExpr) { + auto *AddRec = dyn_cast(NewExpr); + if (!AddRec || !AddRec->isAffine()) + return false; + Expr = AddRec; + return true; +} + +// Recognize monotonic phi variable by matching the following pattern: +// loop_header: +// %monotonic_phi = [%start, %preheader], [%chain_phi0, %latch] +// +// step_bb: +// %step = add/gep %monotonic_phi, %step_val +// +// bbN: +// %chain_phiN = [%monotonic_phi, ], [%step, ] +// +// ... +// +// bb1: +// %chain_phi1 = [%monotonic_phi, ], [%chain_phi2, ] +// +// latch: +// %chain_phi0 = [%monotonic_phi, %pred], [%chain_phi1, %pred] +// +// For this pattern, monotonic phi is described by {%start, +, %step} recurrence +// and predicate is CFG edge %step_bb -> %bbN. +bool MonotonicDescriptor::isMonotonicPHI(PHINode *PN, const Loop *L, + MonotonicDescriptor &Desc, + ScalarEvolution &SE) { + if (!PN->getType()->isIntOrPtrTy() || PN->getParent() != L->getHeader()) + return false; + auto *BackEdgeInst = + dyn_cast(PN->getIncomingValueForBlock(L->getLoopLatch())); + if (!BackEdgeInst) + return false; + PHINode *PHIChain = BackEdgeInst; + std::optional> Inc; + while (PHIChain) { + Desc.Chain.insert(PHIChain); + PHINode *NextPHIChain = nullptr; + for (auto [Block, Incoming] : + zip_equal(PHIChain->blocks(), PHIChain->incoming_values())) { + if (Incoming == PN) + continue; + if (!Incoming->hasOneUse()) + return false; + if (auto *IncomingPHI = dyn_cast(Incoming)) { + if (NextPHIChain) + return false; + NextPHIChain = IncomingPHI; + continue; + } + if (Inc) + return false; + Inc = std::make_pair(Edge{Block, PHIChain->getParent()}, Incoming.get()); + } + PHIChain = NextPHIChain; + } + if (!Inc) + return false; + auto [PredEdge, StepOp] = *Inc; + auto *StepInst = dyn_cast(StepOp); + if (!StepInst) + return false; + Desc.StepInst = StepInst; + Desc.PredEdge = PredEdge; + + // Construct SCEVAddRec for this value. + Value *Start = PN->getIncomingValueForBlock(L->getLoopPreheader()); + + Value *Step = nullptr; + bool StepMatch = + PN->getType()->isPointerTy() + ? match(StepInst, m_PtrAdd(m_Specific(PN), m_Value(Step))) + : match(StepInst, m_Add(m_Specific(PN), m_Value(Step))); + if (!StepMatch || !L->isLoopInvariant(Step)) + return false; + + SCEV::NoWrapFlags WrapFlags = SCEV::FlagAnyWrap; + if (auto *GEP = dyn_cast(StepInst)) { + if (GEP->hasNoUnsignedWrap()) + WrapFlags = ScalarEvolution::setFlags(WrapFlags, SCEV::FlagNUW); + if (GEP->hasNoUnsignedSignedWrap()) + WrapFlags = ScalarEvolution::setFlags(WrapFlags, SCEV::FlagNSW); + } else if (auto *OBO = dyn_cast(StepInst)) { + if (OBO->hasNoUnsignedWrap()) + WrapFlags = ScalarEvolution::setFlags(WrapFlags, SCEV::FlagNUW); + if (OBO->hasNoSignedWrap()) + WrapFlags = ScalarEvolution::setFlags(WrapFlags, SCEV::FlagNSW); + } + + return Desc.setSCEV( + SE.getAddRecExpr(SE.getSCEV(Start), SE.getSCEV(Step), L, WrapFlags)); +} + +bool MonotonicDescriptor::isMonotonicVal(Value *Val, const Loop *L, + MonotonicDescriptor &Desc, + ScalarEvolution &SE) { + if (!Val->getType()->isIntOrPtrTy() || L->isLoopInvariant(Val)) + return false; + auto *CurInst = cast(Val); + + auto LoopVariantVal = [&](Value *V, bool AllowRepeats) { + return L->isLoopInvariant(V) ? nullptr : cast(V); + }; + + while (!isa(CurInst)) { + CurInst = find_singleton(CurInst->operands(), LoopVariantVal); + if (!CurInst) + return false; + }; + + if (!isMonotonicPHI(cast(CurInst), L, Desc, SE)) + return false; + + ValueToSCEVMapTy Map{{CurInst, Desc.getExpr()}}; + return Desc.setSCEV(SCEVParameterRewriter::rewrite(SE.getSCEV(Val), SE, Map)); +} diff --git a/llvm/unittests/Analysis/IVDescriptorsTest.cpp b/llvm/unittests/Analysis/IVDescriptorsTest.cpp index 453800abf9cab..0affe87602b26 100644 --- a/llvm/unittests/Analysis/IVDescriptorsTest.cpp +++ b/llvm/unittests/Analysis/IVDescriptorsTest.cpp @@ -10,6 +10,7 @@ #include "llvm/Analysis/AssumptionCache.h" #include "llvm/Analysis/LoopInfo.h" #include "llvm/Analysis/ScalarEvolution.h" +#include "llvm/Analysis/ScalarEvolutionExpressions.h" #include "llvm/Analysis/TargetLibraryInfo.h" #include "llvm/AsmParser/Parser.h" #include "llvm/IR/Dominators.h" @@ -259,3 +260,155 @@ for.end: EXPECT_EQ(Kind, RecurKind::FMax); }); } + +TEST(IVDescriptorsTest, MonotonicIntVar) { + // Parse the module. + LLVMContext Context; + + std::unique_ptr M = + parseIR(Context, + R"(define void @foo(ptr %dst, i1 %cond, i64 %n) { +entry: + br label %for.body + +for.body: + %i = phi i64 [ 0, %entry ], [ %i.next, %for.inc ] + %monotonic = phi i32 [ 0, %entry ], [ %monotonic.next, %for.inc ] + br i1 %cond, label %if.then, label %for.inc + +if.then: + %inc = add nsw i32 %monotonic, 1 + %monotonic.prom = sext i32 %monotonic to i64 + %arrayidx = getelementptr inbounds i32, ptr %dst, i64 %monotonic.prom + br label %for.inc + +for.inc: + %monotonic.next = phi i32 [ %inc, %if.then ], [ %monotonic, %for.body ] + %i.next = add nuw nsw i64 %i, 1 + %exitcond.not = icmp eq i64 %i.next, %n + br i1 %exitcond.not, label %for.end, label %for.body + +for.end: + ret void +})"); + + runWithLoopInfoAndSE( + *M, "foo", [&](Function &F, LoopInfo &LI, ScalarEvolution &SE) { + Function::iterator FI = F.begin(); + // First basic block is entry - skip it. + BasicBlock *Header = &*(++FI); + assert(Header->getName() == "for.body"); + Loop *L = LI.getLoopFor(Header); + EXPECT_NE(L, nullptr); + BasicBlock::iterator BBI = Header->begin(); + assert((&*BBI)->getName() == "i"); + PHINode *Phi = dyn_cast(&*(++BBI)); + assert(Phi->getName() == "monotonic"); + BasicBlock *IfThen = &*(++FI); + assert(IfThen->getName() == "if.then"); + BBI = IfThen->begin(); + Instruction *StepInst = &*BBI; + assert(StepInst->getName() == "inc"); + Instruction *ExtInst = &*(++BBI); + assert(ExtInst->getName() == "monotonic.prom"); + Instruction *GEPInst = &*(++BBI); + assert(GEPInst->getName() == "arrayidx"); + BasicBlock *IfEnd = &*(++FI); + assert(IfEnd->getName() == "for.inc"); + auto *ChainPhi = dyn_cast(&*(IfEnd->begin())); + assert(ChainPhi->getName() == "monotonic.next"); + // Check %monotonic descriptor. + MonotonicDescriptor Desc; + bool IsMonotonicPhi = + MonotonicDescriptor::isMonotonicPHI(Phi, L, Desc, SE); + EXPECT_TRUE(IsMonotonicPhi); + auto &PhiChain = Desc.getChain(); + EXPECT_TRUE(PhiChain.size() == 1 && PhiChain.contains(ChainPhi)); + EXPECT_EQ(Desc.getStepInst(), StepInst); + EXPECT_EQ(Desc.getPredicateEdge(), + MonotonicDescriptor::Edge(IfThen, IfEnd)); + auto *StartSCEV = SE.getConstant(Phi->getType(), 0); + auto *StepSCEV = SE.getConstant(Phi->getType(), 1); + EXPECT_EQ(Desc.getExpr(), + SE.getAddRecExpr(StartSCEV, StepSCEV, L, SCEV::FlagNW)); + // Check %arrayidx descriptor. + bool IsMonotonicVal = + MonotonicDescriptor::isMonotonicVal(GEPInst, L, Desc, SE); + EXPECT_TRUE(IsMonotonicVal); + // Chain, StepInst and PredicateEdge are the same with %monotonic. + auto &ValChain = Desc.getChain(); + EXPECT_TRUE(ValChain.size() == 1 && ValChain.contains(ChainPhi)); + EXPECT_EQ(Desc.getStepInst(), StepInst); + EXPECT_EQ(Desc.getPredicateEdge(), + MonotonicDescriptor::Edge(IfThen, IfEnd)); + StartSCEV = SE.getSCEV(F.getArg(0)); + StepSCEV = SE.getConstant(StartSCEV->getType(), 4); + EXPECT_EQ(Desc.getExpr(), + SE.getAddRecExpr(StartSCEV, StepSCEV, L, SCEV::FlagNW)); + }); +} + +TEST(IVDescriptorsTest, MonotonicPtrVar) { + // Parse the module. + LLVMContext Context; + + std::unique_ptr M = + parseIR(Context, + R"(define void @foo(ptr %start, i1 %cond, i64 %n) { +entry: + br label %for.body + +for.body: + %i = phi i64 [ 0, %entry ], [ %i.next, %for.inc ] + %monotonic = phi ptr [ %start, %entry ], [ %monotonic.next, %for.inc ] + br i1 %cond, label %if.then, label %for.inc + +if.then: + %inc = getelementptr inbounds i8, ptr %monotonic, i64 4 + br label %for.inc + +for.inc: + %monotonic.next = phi ptr [ %inc, %if.then ], [ %monotonic, %for.body ] + %i.next = add nuw nsw i64 %i, 1 + %exitcond.not = icmp eq i64 %i.next, %n + br i1 %exitcond.not, label %for.end, label %for.body + +for.end: + ret void +})"); + + runWithLoopInfoAndSE( + *M, "foo", [&](Function &F, LoopInfo &LI, ScalarEvolution &SE) { + Function::iterator FI = F.begin(); + // First basic block is entry - skip it. + BasicBlock *Header = &*(++FI); + assert(Header->getName() == "for.body"); + Loop *L = LI.getLoopFor(Header); + EXPECT_NE(L, nullptr); + BasicBlock::iterator BBI = Header->begin(); + assert((&*BBI)->getName() == "i"); + PHINode *Phi = dyn_cast(&*(++BBI)); + assert(Phi->getName() == "monotonic"); + BasicBlock *IfThen = &*(++FI); + assert(IfThen->getName() == "if.then"); + Instruction *StepInst = &*(IfThen->begin()); + assert(StepInst->getName() == "inc"); + BasicBlock *IfEnd = &*(++FI); + assert(IfEnd->getName() == "for.inc"); + auto *ChainPhi = dyn_cast(&*(IfEnd->begin())); + assert(ChainPhi->getName() == "monotonic.next"); + MonotonicDescriptor Desc; + bool IsMonotonicPhi = + MonotonicDescriptor::isMonotonicPHI(Phi, L, Desc, SE); + EXPECT_TRUE(IsMonotonicPhi); + auto &Chain = Desc.getChain(); + EXPECT_TRUE(Chain.size() == 1 && Chain.contains(ChainPhi)); + EXPECT_EQ(Desc.getStepInst(), StepInst); + EXPECT_EQ(Desc.getPredicateEdge(), + MonotonicDescriptor::Edge(IfThen, IfEnd)); + auto *StartSCEV = SE.getSCEV(F.getArg(0)); + auto *StepSCEV = SE.getConstant(StartSCEV->getType(), 4); + EXPECT_EQ(Desc.getExpr(), + SE.getAddRecExpr(StartSCEV, StepSCEV, L, SCEV::FlagNW)); + }); +}