[Mlir-commits] [mlir] [mlir][IntRangeAnalysis] Fix assertion in inferAffineExpr for mod with range crossing modulus boundary (PR #188842)

[Zhewen Yu]

https://github.com/Yu-Zhewen created https://github.com/llvm/llvm-project/pull/188842

The "small range with constant divisor" optimization in inferAffineExpr for AffineExprKind::Mod assumed that if the dividend range span (lhsMax - lhsMin) is less than the divisor, then the mod results form a contiguous range. This is not always true, as the range can straddle a modulus boundary.

For example, [14, 17] mod 8:
- Span is 3 < 8, so the old condition passed
- But 14%8=6 and 17%8=1 (wraps at 16)
- umin=6, umax=1 → assertion umin.ule(umax) fails

The fix adds a same-quotient check (lhsMin/rhs == lhsMax/rhs) to ensure both endpoints fall within the same modular period. When they don't, we fall back to the conservative [0, divisor-1] range.

Assisted-by: Cursor (Claude)


>From 1a914a86be232614fa8bb43b68d7aa1a1028b386 Mon Sep 17 00:00:00 2001
From: Yu-Zhewen <zhewenyu at amd.com>
Date: Thu, 26 Mar 2026 20:57:24 +0000
Subject: [PATCH] fix bound

Signed-off-by: Yu-Zhewen <zhewenyu at amd.com>
---
 mlir/lib/Interfaces/Utils/InferIntRangeCommon.cpp | 15 +++++++--------
 mlir/test/Dialect/Affine/int-range-interface.mlir | 12 ++++++++++++
 2 files changed, 19 insertions(+), 8 deletions(-)

diff --git a/mlir/lib/Interfaces/Utils/InferIntRangeCommon.cpp b/mlir/lib/Interfaces/Utils/InferIntRangeCommon.cpp
index 21f07ddce4495..067000116a46e 100644
--- a/mlir/lib/Interfaces/Utils/InferIntRangeCommon.cpp
+++ b/mlir/lib/Interfaces/Utils/InferIntRangeCommon.cpp
@@ -854,17 +854,16 @@ mlir::intrange::inferAffineExpr(AffineExpr expr,
       umax = lhsMax;
     }
     // Special case: sweeping out a contiguous range with constant divisor.
-    // Only applies when dividend is non-negative to ensure result range is
-    // contiguous.
+    // Only applies when dividend is non-negative and the range does not
+    // cross a modulus boundary (same quotient), ensuring contiguity.
     else if (rhsMin == rhsMax && lhsMin.isNonNegative() &&
-             (lhsMax - lhsMin).ult(rhsMax)) {
-      // For non-negative dividends, Euclidean mod is same as unsigned
-      // remainder.
+             (lhsMax - lhsMin).ult(rhsMax) &&
+             lhsMin.udiv(rhsMax) == lhsMax.udiv(rhsMax)) {
+      // For non-negative dividends within the same modular period,
+      // Euclidean mod is same as unsigned remainder and the result is
+      // contiguous.
       umin = lhsMin.urem(rhsMax);
       umax = lhsMax.urem(rhsMax);
-      // Result should be contiguous since we're not wrapping around.
-      assert(umin.ule(umax) &&
-             "Range should be contiguous for non-negative dividend");
     }
 
     return ConstantIntRanges::fromUnsigned(umin, umax);
diff --git a/mlir/test/Dialect/Affine/int-range-interface.mlir b/mlir/test/Dialect/Affine/int-range-interface.mlir
index ca8df8a9e5788..ac64ad09ee244 100644
--- a/mlir/test/Dialect/Affine/int-range-interface.mlir
+++ b/mlir/test/Dialect/Affine/int-range-interface.mlir
@@ -147,6 +147,18 @@ func.func @affine_apply_mod_variable_divisor() -> index {
   func.return %1 : index
 }
 
+// CHECK-LABEL: func @affine_apply_mod_cross_boundary
+// CHECK: test.reflect_bounds {smax = 7 : index, smin = 0 : index, umax = 7 : index, umin = 0 : index}
+func.func @affine_apply_mod_cross_boundary() -> index {
+  %d0 = test.with_bounds { umin = 14 : index, umax = 17 : index,
+                           smin = 14 : index, smax = 17 : index } : index
+  // Range [14, 17] spans a mod-8 boundary (at 16): 14%8=6, 17%8=1.
+  // Span 3 < 8 but the range wraps, so we fall back to [0, 7].
+  %0 = affine.apply affine_map<(d0) -> (d0 mod 8)>(%d0)
+  %1 = test.reflect_bounds %0 : index
+  func.return %1 : index
+}
+
 // CHECK-LABEL: func @affine_apply_mod_negative_dividend
 // CHECK: test.reflect_bounds {smax = 3 : index, smin = 0 : index, umax = 3 : index, umin = 0 : index}
 func.func @affine_apply_mod_negative_dividend() -> index {