[IVDescriptors] Implement MonotonicDescriptor

Author: skachkov-sc

llvm/include/llvm/Analysis/IVDescriptors.h

@@ -27,6 +27,7 @@ class Loop;
 class PredicatedScalarEvolution;
 class ScalarEvolution;
 class SCEV;
+class SCEVAddRecExpr;
 class StoreInst;
 
 /// These are the kinds of recurrences that we support.
@@ -426,6 +427,44 @@ class InductionDescriptor {
   SmallVector<Instruction *, 2> 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<BasicBlock *, BasicBlock *>;
+
+  MonotonicDescriptor() = default;
+
+  const SmallPtrSetImpl<PHINode *> &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<PHINode *, 1> Chain;
+  Instruction *StepInst;
+  Edge PredEdge;
+  const SCEVAddRecExpr *Expr;
+
+  bool setSCEV(const SCEV *NewExpr);
+};
+
 } // end namespace llvm
 
 #endif // LLVM_ANALYSIS_IVDESCRIPTORS_H

llvm/lib/Analysis/IVDescriptors.cpp

@@ -1590,3 +1590,122 @@ bool InductionDescriptor::isInductionPHI(
   D = InductionDescriptor(StartValue, IK_PtrInduction, Step);
   return true;
 }
+
+bool MonotonicDescriptor::setSCEV(const SCEV *NewExpr) {
+  auto *AddRec = dyn_cast<SCEVAddRecExpr>(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<PHINode>(PN->getIncomingValueForBlock(L->getLoopLatch()));
+  if (!BackEdgeInst)
+    return false;
+  SmallVector<PHINode *, 4> Worklist{BackEdgeInst};
+  std::optional<std::pair<Edge, Value *>> Inc;
+  while (!Worklist.empty()) {
+    auto *Phi = Worklist.pop_back_val();
+    Desc.Chain.insert(Phi);
+    for (unsigned I = 0, E = Phi->getNumOperands(); I != E; ++I) {
+      auto *IncomingVal = Phi->getIncomingValue(I);
+      if (IncomingVal == PN)
+        continue;
+      if (!IncomingVal->hasOneUse())
+        return false;
+      if (auto *IncomingPhi = dyn_cast<PHINode>(IncomingVal)) {
+        Worklist.push_back(IncomingPhi);
+        continue;
+      }
+      if (Inc)
+        return false;
+      Inc = std::make_pair(Edge{Phi->getIncomingBlock(I), Phi->getParent()},
+                           IncomingVal);
+    }
+  }
+  if (!Inc)
+    return false;
+  auto [PredEdge, StepOp] = *Inc;
+  auto *StepInst = dyn_cast<Instruction>(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<GEPOperator>(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<OverflowingBinaryOperator>(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<Instruction>(Val);
+
+  auto NonInvariantVal = [&](Value *V, bool AllowRepeats) {
+    return L->isLoopInvariant(V) ? nullptr : cast<Instruction>(V);
+  };
+
+  while (!isa<PHINode>(CurInst)) {
+    CurInst = find_singleton<Instruction>(CurInst->operands(), NonInvariantVal);
+    if (!CurInst)
+      return false;
+  };
+
+  if (!isMonotonicPHI(cast<PHINode>(CurInst), L, Desc, SE))
+    return false;
+
+  ValueToSCEVMapTy Map{{CurInst, Desc.getExpr()}};
+  return Desc.setSCEV(SCEVParameterRewriter::rewrite(SE.getSCEV(Val), SE, Map));
+}