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

[Kiran Chandramohan via llvm-commits]

================
@@ -737,6 +752,367 @@ 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,
+/// MulExpr has the constant operand first.
+/// Returns std::nullopt if the expression contains non-stencil terms.
+static std::optional<StencilDecomposition>
+decomposeStencilOffset(const SCEV *Expr, ScalarEvolution &SE, const Loop &L) {
+  StencilDecomposition D;
+
+  // Collect top-level additive terms.
+  SmallVector<const SCEV *, 4> Terms;
+  if (auto *Add = dyn_cast<SCEVAddExpr>(Expr))
+    Terms.append(Add->operands().begin(), Add->operands().end());
+  else
+    Terms.push_back(Expr);
+
+  for (const SCEV *Term : Terms) {
+    if (auto *C = dyn_cast<SCEVConstant>(Term)) {
+      D.Constant += C->getAPInt().getSExtValue();
+    } else if (auto *Mul = dyn_cast<SCEVMulExpr>(Term)) {
+      // SCEV canonical form: constant is operand 0 in a MulExpr.
+      if (Mul->getNumOperands() != 2)
+        return std::nullopt;
+      auto *C = dyn_cast<SCEVConstant>(Mul->getOperand(0));
+      if (!C || !SE.isLoopInvariant(Mul->getOperand(1), &L))
+        return std::nullopt;
+      D.Coefficients[Mul->getOperand(1)] += C->getAPInt().getSExtValue();
+    } else if (SE.isLoopInvariant(Term, &L)) {
+      D.Coefficients[Term] += 1;
+    } else {
+      return std::nullopt;
+    }
+  }
+  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;
+
+  // Count the total checks from current grouping.
+  unsigned TotalChecks = 0;
+  for (unsigned I = 0; I < CheckingGroups.size(); ++I)
+    for (unsigned J = I + 1; J < CheckingGroups.size(); ++J)
+      if (needsChecking(CheckingGroups[I], CheckingGroups[J]))
+        ++TotalChecks;
+
+  if (!EnableStencilMerge && TotalChecks <= StencilMergeCheckThreshold) {
----------------
kiranchandramohan wrote:

Do you intend to enable the pass when (TotalChecks > StencilMergeCheckThreshold) even when the Flag (EnableStencilMerge) is OFF?

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