[Mlir-commits] [mlir] 3bfed00 - [MLIR] Guard against overflow when simplifying affine expressions. (#228010)

[llvmlistbot at llvm.org]

Author: Jiaqi He
Date: 2026-10-12T12:38:03+09:00
New Revision: 3bfed000ae79c2fdb6785e2b1e2222ba1db6d87a

URL: https://github.com/llvm/llvm-project/commit/3bfed000ae79c2fdb6785e2b1e2222ba1db6d87a
DIFF: https://github.com/llvm/llvm-project/commit/3bfed000ae79c2fdb6785e2b1e2222ba1db6d87a.diff

LOG: [MLIR] Guard against overflow when simplifying affine expressions. (#228010)

Fixes https://github.com/llvm/llvm-project/issues/224577.

Current canonicalizer may lead to signed integer overflow while
simplifying affine expressions. This change adds overflow checks and
returns `failure()` when an overflow is detected.

Added: 
    

Modified: 
    mlir/lib/IR/AffineExpr.cpp
    mlir/test/Dialect/Affine/canonicalize.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/IR/AffineExpr.cpp b/mlir/lib/IR/AffineExpr.cpp
index 05d643d36896c..dd9a1f588a842 100644
--- a/mlir/lib/IR/AffineExpr.cpp
+++ b/mlir/lib/IR/AffineExpr.cpp
@@ -1356,8 +1356,12 @@ LogicalResult SimpleAffineExprFlattener::visitMulExpr(AffineBinaryOpExpr expr) {
 
   // Get the RHS constant.
   int64_t rhsConst = rhs[getConstantIndex()];
-  for (int64_t &lhsElt : lhs)
-    lhsElt *= rhsConst;
+  for (int64_t &lhsElt : lhs) {
+    auto [prod, overflow] = llvm::MulOverflow(lhsElt, rhsConst);
+    if (overflow)
+      return failure();
+    lhsElt = prod;
+  }
 
   return success();
 }
@@ -1369,7 +1373,10 @@ LogicalResult SimpleAffineExprFlattener::visitAddExpr(AffineBinaryOpExpr expr) {
   assert(lhs.size() == rhs.size());
   // Update the LHS in place.
   for (unsigned i = 0, e = rhs.size(); i < e; i++) {
-    lhs[i] += rhs[i];
+    auto [sum, overflow] = llvm::AddOverflow(lhs[i], rhs[i]);
+    if (overflow)
+      return failure();
+    lhs[i] = sum;
   }
   // Pop off the RHS.
   operandExprStack.pop_back();

diff  --git a/mlir/test/Dialect/Affine/canonicalize.mlir b/mlir/test/Dialect/Affine/canonicalize.mlir
index 57b92e79403fe..d899dd0772d4a 100644
--- a/mlir/test/Dialect/Affine/canonicalize.mlir
+++ b/mlir/test/Dialect/Affine/canonicalize.mlir
@@ -2653,3 +2653,39 @@ func.func @compose_into_access_keeps_alignment(%memref: memref<100xi32>, %i: ind
   affine.store %val, %memref[%idx] { alignment = 16 } : memref<100xi32>
   return
 }
+
+// -----
+
+// Check that canonicalizing an affine expression with addition overflow
+// does not crash.
+
+// CHECK-LABEL: func.func @index_boundary_add
+// CHECK: affine.load %{{.*}}[-9223372036854775808, -9223372036854775808 - 1] : memref<64x64xf32>
+
+func.func @index_boundary_add() -> f32 {
+  %large = arith.constant 9223372036854775807 : i64
+  %idx = arith.index_cast %large : i64 to index
+  %c1 = arith.constant 1 : index
+  %r = arith.addi %idx, %c1 : index
+  %A = memref.alloc() : memref<64x64xf32>
+  %0 = affine.load %A[%r, %r - 1] : memref<64x64xf32>
+  return %0 : f32
+}
+
+// -----
+
+// Check that canonicalizing an affine expression with multiplication overflow
+// does not crash.
+
+// CHECK-LABEL: func.func @index_boundary_mul
+// CHECK: affine.load %{{.*}}[-9223372036854775808, -9223372036854775808 * 2] : memref<64x64xf32>
+
+func.func @index_boundary_mul() -> f32 {
+  %large = arith.constant 9223372036854775807 : i64
+  %idx = arith.index_cast %large : i64 to index
+  %c1 = arith.constant 1 : index
+  %r = arith.addi %idx, %c1 : index
+  %A = memref.alloc() : memref<64x64xf32>
+  %0 = affine.load %A[%r, %r * 2] : memref<64x64xf32>
+  return %0 : f32
+}