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

[David Sherwood via llvm-commits]

================
@@ -737,6 +763,415 @@ void RuntimePointerChecking::groupChecks(
   }
 }
 
+/// Result of decomposing a SCEV expression into stencil offset form:
+///   Offset = Constant + sum(Coefficients[stride] * stride)
+/// where each stride is a loop-invariant SCEV expression.
+struct StencilDecomposition {
+  int64_t Constant = 0;
+  /// Map from loop-invariant stride SCEV to its integer coefficient.
+  SmallMapVector<const SCEV *, int64_t, 4> Coefficients;
+};
+
+/// Try to decompose \p Expr into a stencil offset function of loop-invariant
+/// strides: C + a1*s1 + a2*s2 + ...
+/// Relies on SCEV's canonical form: AddExpr operands are flattened (N-ary),
+/// MulExpr has the constant operand first when present.
+/// Returns std::nullopt if the expression contains non-stencil terms or any
+/// SCEV constant doesn't fit in int64_t (we commit to the signed
+/// interpretation; values that need more than 64 significant bits are
+/// out of scope).
+static std::optional<StencilDecomposition>
+decomposeStencilOffset(const SCEV *Expr, ScalarEvolution &SE, const Loop &L) {
+  StencilDecomposition D;
+
+  // The only caller passes a difference of two pointer Starts, and a Start is
+  // always loop-invariant (getStartAndEndForAccess asserts it). So Expr is
+  // loop-invariant and every term below is loop-invariant too.
+  assert(SE.isLoopInvariant(Expr, &L) && "expected a loop-invariant offset");
+
+  // Collect top-level additive terms.
+  SmallVector<const SCEV *, 4> Terms;
+  if (auto *Add = dyn_cast<SCEVAddExpr>(Expr))
+    append_range(Terms, Add->operands());
+  else
+    Terms.push_back(Expr);
+
+  const SCEVConstant *C;
+  const SCEV *Stride;
+  for (const SCEV *Term : Terms) {
+    if (match(Term, m_SCEVConstant(C))) {
+      auto V = C->getAPInt().trySExtValue();
+      if (!V)
+        return std::nullopt;
+      D.Constant += *V;
+    } else if (match(Term, m_scev_Mul(m_SCEVConstant(C), m_SCEV(Stride)))) {
+      assert(SE.isLoopInvariant(Stride, &L) && "stride must be loop-invariant");
+      auto V = C->getAPInt().trySExtValue();
+      if (!V)
+        return std::nullopt;
+      D.Coefficients[Stride] += *V;
+    } else {
+      assert(SE.isLoopInvariant(Term, &L) && "term must be loop-invariant");
+      D.Coefficients[Term] += 1;
+    }
+  }
+  return D;
+}
+
+void RuntimePointerChecking::mergeStencilGroups(PredicatedScalarEvolution &PSE,
+                                                Loop &L) {
+  LLVM_DEBUG(dbgs() << "LAA: Attempting stencil group merging on "
+                    << CheckingGroups.size() << " groups\n");
+
+  if (CheckingGroups.size() < 2)
+    return;
+
+  // We try to merge groups produced by groupChecks when their pointers follow
+  // a stencil access pattern. groupChecks intentionally only merges pointers
+  // whose min/max bounds differ by constants, because that keeps each runtime
+  // check precise. Stencil kernels often read the same underlying object at
+  // several loop-invariant stride offsets, so those groups remain separate and
+  // can produce too many checks. Here we trade some precision for fewer
+  // checks by replacing a set of same-dependence, same-alias read groups with
+  // one conservative bounding group.
----------------
david-arm wrote:

This is great! Thanks for adding these comments @igogo-x86. I find the high-level explanation really helps.

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