[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.
+  //
+  // We use the following algorithm to construct a merged stencil group:
+  //   - collect checking groups that share both DependencySetId and AliasSetId;
+  //   - reject groups with writes, predicated accesses, different access
+  //     ranges, or different recurrence steps;
+  //   - use one member as the base and decompose each other member's offset
+  //     from that base as C + sum(Coeff[Stride] * Stride), where Stride is
+  //     loop-invariant;
+  //   - build one bounding range by taking the min/max constant offset and the
+  //     min/max coefficient for each stride, adding predicates for strides
+  //     that are not already known positive;
+  //   - commit the merge only if the local cost model reduces the number of
+  //     checks after accounting for any new predicates.
+
+  // Stencil merging runs when either:
+  //   - the flag is set to 'force' (-stencil-runtime-check-merge=force), or
+  //   - the flag is set to 'auto' (-stencil-runtime-check-merge=auto) AND the
+  //     current check count exceeds the auto-trigger threshold, where the
----------------
david-arm wrote:

I think this isn't quite true, because the metric you're using here is different to the one used by the loop vectoriser. It could be that the loop vectoriser wouldn't reject the original checks *and* we still decide to do a stencil merge. In theory, this could be worse than before. I'm fine with that, but it might be worth saying that we're using an approximation for the loop vectoriser's cutoff.

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