[IVDescriptors] Implement MonotonicDescriptor

[skachkov-sc]

RFC link: https://discourse.llvm.org/t/rfc-loop-vectorization-of-compress-store-expand-load-patterns/86442

"Monotonic" variable is similar to induction variable, but its value is updated under some condition, e.g.:

int idx = 0;
for(int i = 0; i < n; ++i) {
  // some uses of idx
  if (cond)
    ++idx;
}

In this example, i is induction variable and idx is monotonic variable: it's updated only when cond == true. In LLVM IR, this looks like:

loop_header:
  %monotonic_phi = [%start, %prehader], [ %chain_phi0, %latch]

step_bb:
  %step = add/gep %monotonic_phi, %step_val

bbN:
  %chain_phiN = [%step, %step_bb], [%monotonic_phi, %some_pred]

...

bb1:
  %chain_phi1 = [%monotonic_phi, %some_pred], [%chain_phi2, %some_pred]

loop_latch:
  %chain_phi0 = [%monotonic_phi, %some_pred], [%chain_phi1, %some_pred]

We start the analysis from the header phi (%monotonic_phi). Its backedge value (%chain_phi0) can be either unmodified (%monotonic_val) or the incremented one (%step) depending on some condition. Therefore, all incoming values in chain phi should be %monotonic_phi except the one that points to the next chain phi or the step instruction. The analysis stops when the correct step instruction is found (correct step instruction is in form add/gep %monotonic_phi, %step_val, and %step_val is loop-invariant). In this way, %monotonic_val is incremented on %step_val only when last chain phi "selects" %step incoming value; in other words, when step_bb -> bbN edge is executed. The condition of monotonic variable update is stored as such CFG edge (getPredicateEdge() method). The other MonotonicDescriptor methods are:

1. getChain() - returns set of chain phis (%chain_phi0, ..., %chain_phiN)
2. getStepInst() - returns %step instruction (either add or getlementptr)
3. getExpr() - returns SCEVAddRec in assumption that condition of update is always true (SCEVAddRec + condtion predicate fully describes evolution of monotonic variable in the loop)

MonotonicDescriptor also implements isMonotonicVal method. This routine will allow to recognize values which have monotonic phi as a transitive dependency, e.g. array subscript operators arr[idx], where idx is monotonic phi. The descriptor of monotonic value is equivalent to the descriptor of monotonic phi, except its SCEV: we obtain the SCEV of monotonic value and replace the SCEVUnknown expression that corresponds to monotonic phi with the SCEVAddRec from monotonic phi descriptor. The resulted SCEV shows the evolution of monotonic value in loop iterations when predicate is true (similarly to SCEV).

Restrictions

1. We only support "post-increment" update of monotonic vars; in other words, uses of chain phis (other than in other chain phis) are prohibited. It's only alowed to use %monotonic_phi outside of described chain pattern.
2. We don't support select instructions in monotonic patterns (instead of chain phis). In theory this can be implemented, but on practice we want to use MonotonicDescriptor to recognize expandloads/compressstores, where step instruction is placed near memory instruction, so if-conversion to select is not applied there (because basic block can't be eliminated anyway). So, this restriction doesn't limit our abilities to recognize expandloads/compressstore patterns.

[llvmbot]

@llvm/pr-subscribers-llvm-analysis

Author: Sergey Kachkov (skachkov-sc)

Changes

---

Full diff: https://github.com/llvm/llvm-project/pull/140720.diff

2 Files Affected:

• (modified) llvm/include/llvm/Analysis/IVDescriptors.h (+39)
• (modified) llvm/lib/Analysis/IVDescriptors.cpp (+119)

diff --git a/llvm/include/llvm/Analysis/IVDescriptors.h b/llvm/include/llvm/Analysis/IVDescriptors.h
index 140edff13a67f..eedb7eee80fad 100644
--- a/llvm/include/llvm/Analysis/IVDescriptors.h
+++ b/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
diff --git a/llvm/lib/Analysis/IVDescriptors.cpp b/llvm/lib/Analysis/IVDescriptors.cpp
index a273338670164..1818fe4d16e9a 100644
--- a/llvm/lib/Analysis/IVDescriptors.cpp
+++ b/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));
+}

[skachkov-sc]

Gentle ping (this patch is NFC on its own, but is required for #140721 and #140723)

[MacDue]

Is there a way this code can be tested in isolation? This PR is adding functional changes; it's just unused/untested within this PR.

[skachkov-sc]

Is there a way this code can be tested in isolation? This PR is adding function changes; it's just unused/untested within this PR.

I think we can only make unit tests there; added tests for MonotonicDescriptor to check recognition of integer/pointer monotonic phis.

[fhahn]

Could you please add more detailed descriptions to this and the related PRs, making it a bit easier to review?

[MacDue]

nit: Avoid a double-negative

Suggested change

	auto NonInvariantVal = [&](Value *V, bool AllowRepeats) {
	auto LoopVariantVal = [&](Value *V, bool AllowRepeats) {

[MacDue]

Since these two tests are so similar, is it worth testing isMonotonicVal on inc in this case? So both methods are covered.

[skachkov-sc]

Added testing of isMonotonicVal() in MonotonicIntVar; this should cover the case we are mostly interested in - when loop has integer monotonic phi which is used as an index in array subscript operator (and we should recognize the access to this array as "monotonic"; when monotonic access has the stride equals of element size that means that the access is "compressed")

[MacDue]

IIUC, only the only a single the monotonic PHI is the only loop-variant value allowed to be used in transitively?

[skachkov-sc]

Yes, we allow only single transitive use of monotonic phi here (added paragraph about isMonotonicVal() in the PR's description). I can't recall any problems with multiple uses, but such cases are not very common in practice (the main intention of this method is to recognize instruction sequences like MonotonicPhi -> [sext/zext] -> GEP)