[llvm] [LAA] Add stencil group merging to reduce runtime pointer checks (PR #187252)

[David Sherwood via llvm-commits]

================
@@ -989,6 +989,95 @@ getStencilStrideUpperLimit(const StencilDecomposition &D, unsigned BitWidth) {
   return APInt(BitWidth, CoeffSum ? Budget / CoeffSum : SignedMax);
 }
 
+namespace {
+/// The runtime checks the merge needs on each stride.
+/// Example:
+///   s1: {NeedsPositive = true, Max = 1000} means the checks 1 <= s1 <= 1000
+///   s2: {Max = 50} means the check s2 <= 50
+/// The lower limit is always 1, so a flag is enough for it.
+/// Several members can each ask for an upper limit on the same stride, but only
+/// the smallest one is kept.
+class StrideLimits {
+  struct Limit {
+    bool NeedsPositive = false;
+    std::optional<APInt> Max;
+  };
+  SmallMapVector<const SCEV *, Limit, 4> Limits;
+
+public:
+  void requireLowerLimit(const SCEV *Stride) {
+    Limits[Stride].NeedsPositive = true;
+  }
+
+  void requireUpperLimit(const SCEV *Stride, const APInt &Max) {
+    std::optional<APInt> &Current = Limits[Stride].Max;
+    if (!Current || Max.ult(*Current))
+      Current = Max;
+  }
+
+  /// Add every new or more strict check in \p Other to this set.
+  void addFrom(const StrideLimits &Other) {
+    for (const auto &[Stride, L] : Other.Limits) {
+      if (L.NeedsPositive)
+        requireLowerLimit(Stride);
+      if (L.Max)
+        requireUpperLimit(Stride, *L.Max);
+    }
+  }
+
+  /// Count the strides that have no check in \p Committed yet.
+  unsigned countNew(const StrideLimits &Committed) const {
+    return count_if(Limits, [&](const auto &Entry) {
+      return !Committed.Limits.contains(Entry.first);
+    });
+  }
+
+  /// Add the checks to \p PSE as SCEV predicates.
+  void addPredicates(PredicatedScalarEvolution &PSE) const {
+    ScalarEvolution &SE = *PSE.getSE();
+    for (const auto &[Stride, L] : Limits) {
+      if (!L.NeedsPositive)
----------------
david-arm wrote:

OK, that's a bit confusing. Why not just use a single loop? For example,

```
  for (const auto &[Stride, L] : Limits) {
    if (!L.NeedsPositive && !L.Max)
      continue
    if (L.NeedsPositive) {
      const SCEV *Zero = SE.getZero(Stride->getType());
      PSE.addPredicate(
          *SE.getComparePredicate(ICmpInst::ICMP_SGT, Stride, Zero));
      LLVM_DEBUG(dbgs() << "LAA:   Adding positive-stride predicate for "
                        << *Stride << "\n");
    }
    if (L.Max) {
      PSE.addPredicate(*SE.getComparePredicate(ICmpInst::ICMP_SLE, Stride,
                                               SE.getConstant(*L.Max)));
      LLVM_DEBUG(dbgs() << "LAA:   Adding stride upper-limit predicate "
                        << *Stride << " <= " << *L.Max << "\n");
    }
  }
```

which is a bit easier to follow (for me at least) and more efficient? Is there a reason why you prefer to group positive-stride and upper-limit predicates together?

https://github.com/llvm/llvm-project/pull/187252