[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 updated https://github.com/llvm/llvm-project/pull/188842

>From 83101510ea5dbf97bc893a780cb47a5e7414418f 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>
Made-with: Cursor
Signed-off-by: Yu-Zhewen <zhewenyu at amd.com>
Made-with: Cursor
---
 mlir/lib/Interfaces/Utils/InferIntRangeCommon.cpp | 12 +++++++-----
 mlir/test/Dialect/Affine/int-range-interface.mlir | 12 ++++++++++++
 2 files changed, 19 insertions(+), 5 deletions(-)

diff --git a/mlir/lib/Interfaces/Utils/InferIntRangeCommon.cpp b/mlir/lib/Interfaces/Utils/InferIntRangeCommon.cpp
index 21f07ddce4495..c9f49fda726e7 100644
--- a/mlir/lib/Interfaces/Utils/InferIntRangeCommon.cpp
+++ b/mlir/lib/Interfaces/Utils/InferIntRangeCommon.cpp
@@ -854,12 +854,14 @@ 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.
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 {