[Mlir-commits] [mlir] [mlir][affine] Add useExpensiveMath option to AffineLoopNormalize pass (PR #211989)

[lonely eagle]

================
@@ -28,13 +30,79 @@ using namespace mlir::affine;
 
 namespace {
 
+/// Computes the constant upper or lower bound for a given affine map expression
+/// and its operands, constrained by the specified type.
+static FailureOr<int64_t> computeConstantBound(AffineMap map,
+                                               ValueRange operands,
+                                               presburger::BoundType type) {
+  ValueBoundsConstraintSet::Variable var(map, operands);
+  return ValueBoundsConstraintSet::computeConstantBound(
+      type, var, nullptr, {/*closedUb=*/true, /*allowIntegerType=*/true});
+}
+
+/// Attempts to infer a static constant upper bound for the given normalized
+/// `affine.for` loop using Value Bounds Analysis. If the dynamic upper bound's
+/// range [upperMin, upperMax] is proven to be a single constant value (upperMin
+/// == upperMax), the upper bound is directly replaced with this constant.
+/// Otherwise, if upperMin > 0, the loop is split (peeled) into a static main
+/// loop with a constant upper bound (`upperMin`) and a residual tail loop
+/// iterating from `upperMin` to the original dynamic bound.
+static LogicalResult
+inferAffineLoopUpperConstantBound(RewriterBase &b, AffineForOp forOp,
+                                  bool promoteSingleIter = true) {
+  // Ensure the loop is normalized (lower bound is strictly 0 and step is 1).
+  if (!forOp.hasConstantLowerBound() || forOp.getConstantLowerBound() != 0)
+    return failure();
+  if (forOp.getStepAsInt() != 1)
+    return failure();
+  if (forOp.getUpperBoundMap().getNumResults() > 1)
+    return failure();
+
+  // Infer the range [upperMin, upperMax] for the upper bound. We require a
+  // strictly positive minimum bound (upperMin > 0) to guarantee a safe,
+  // non-empty static trip count for the main loop.
+  FailureOr<int64_t> upperMin = computeConstantBound(
+      forOp.getUpperBoundMap(), forOp.getUpperBoundOperands(),
+      presburger::BoundType::LB);
+  FailureOr<int64_t> upperMax = computeConstantBound(
+      forOp.getUpperBoundMap(), forOp.getUpperBoundOperands(),
+      presburger::BoundType::UB);
+  if (failed(upperMin) || *upperMin <= 0)
+    return failure();
+
+  // The upper bound is dynamic within [upperMin, upperMax]. Split the loop into
+  // a static main loop (0 to upperMin) and a residual tail loop (upperMin to
+  // dynamic bound).
+  if (failed(upperMax) || *upperMax > *upperMin) {
+    b.setInsertionPoint(forOp);
+    AffineForOp clonedForOp = cast<AffineForOp>(b.clone(*forOp));
+    clonedForOp.setConstantUpperBound(*upperMin);
----------------
linuxlonelyeagle wrote:

I double-checked the code and realized I misunderstood—clonedForOp is actually the op where we set the constant upper bound.

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