[Mlir-commits] [mlir] [mlir][affine] avoid folding affine expressions that overflow int64_t (PR #231085)
Mahmoud Naderi
llvmlistbot at llvm.org
Tue Oct 13 02:10:35 PDT 2026
https://github.com/slant14 updated https://github.com/llvm/llvm-project/pull/231085
>From 8da7047c364e095402ddfb943e14c21529cc1b9b Mon Sep 17 00:00:00 2001
From: Mahmoud <mahmoud at Mahmouds-MacBook-Pro.local>
Date: Mon, 12 Oct 2026 09:51:16 +0300
Subject: [PATCH 1/2] [mlir][affine] avoid folding affine expressions that
overflow int64_t
---
mlir/lib/IR/AffineExpr.cpp | 10 ++--
mlir/lib/IR/AffineMap.cpp | 23 +++++++--
mlir/test/Dialect/Affine/canonicalize.mlir | 57 ++++++++++++++++++++++
3 files changed, 82 insertions(+), 8 deletions(-)
diff --git a/mlir/lib/IR/AffineExpr.cpp b/mlir/lib/IR/AffineExpr.cpp
index 05d643d36896c..209ad23f980b2 100644
--- a/mlir/lib/IR/AffineExpr.cpp
+++ b/mlir/lib/IR/AffineExpr.cpp
@@ -1354,10 +1354,11 @@ LogicalResult SimpleAffineExprFlattener::visitMulExpr(AffineBinaryOpExpr expr) {
return addLocalVariableSemiAffine(mulLhs, rhs, a * b, lhs, lhs.size());
}
- // Get the RHS constant.
+ // Get the RHS constant. Bail out if any coefficient overflows.
int64_t rhsConst = rhs[getConstantIndex()];
for (int64_t &lhsElt : lhs)
- lhsElt *= rhsConst;
+ if (llvm::MulOverflow(lhsElt, rhsConst, lhsElt))
+ return failure();
return success();
}
@@ -1367,9 +1368,10 @@ LogicalResult SimpleAffineExprFlattener::visitAddExpr(AffineBinaryOpExpr expr) {
const auto &rhs = operandExprStack.back();
auto &lhs = operandExprStack[operandExprStack.size() - 2];
assert(lhs.size() == rhs.size());
- // Update the LHS in place.
+ // Update the LHS in place. Bail out if any coefficient overflows.
for (unsigned i = 0, e = rhs.size(); i < e; i++) {
- lhs[i] += rhs[i];
+ if (llvm::AddOverflow(lhs[i], rhs[i], lhs[i]))
+ return failure();
}
// Pop off the RHS.
operandExprStack.pop_back();
diff --git a/mlir/lib/IR/AffineMap.cpp b/mlir/lib/IR/AffineMap.cpp
index 446bc9a6f6f87..dfbf43392e8af 100644
--- a/mlir/lib/IR/AffineMap.cpp
+++ b/mlir/lib/IR/AffineMap.cpp
@@ -24,14 +24,15 @@ using namespace mlir;
using llvm::divideCeilSigned;
using llvm::divideFloorSigned;
+using llvm::divideSignedWouldOverflow;
using llvm::mod;
namespace {
-
// AffineExprConstantFolder evaluates an affine expression using constant
// operands passed in 'operandConsts'. Returns an IntegerAttr attribute
// representing the constant value of the affine expression evaluated on
-// constant 'operandConsts', or nullptr if it can't be folded.
+// constant 'operandConsts', or nullptr if it can't be folded (e.g., because the
+// result would overflow int64_t).
class AffineExprConstantFolder {
public:
AffineExprConstantFolder(unsigned numDims, ArrayRef<Attribute> operandConsts)
@@ -52,10 +53,20 @@ class AffineExprConstantFolder {
switch (expr.getKind()) {
case AffineExprKind::Add:
return constantFoldBinExpr(
- expr, [](int64_t lhs, int64_t rhs) { return lhs + rhs; });
+ expr, [](int64_t lhs, int64_t rhs) -> std::optional<int64_t> {
+ int64_t sum;
+ if (llvm::AddOverflow(lhs, rhs, sum))
+ return std::nullopt;
+ return sum;
+ });
case AffineExprKind::Mul:
return constantFoldBinExpr(
- expr, [](int64_t lhs, int64_t rhs) { return lhs * rhs; });
+ expr, [](int64_t lhs, int64_t rhs) -> std::optional<int64_t> {
+ int64_t product;
+ if (llvm::MulOverflow(lhs, rhs, product))
+ return std::nullopt;
+ return product;
+ });
case AffineExprKind::Mod:
return constantFoldBinExpr(
expr, [this](int64_t lhs, int64_t rhs) -> std::optional<int64_t> {
@@ -72,6 +83,8 @@ class AffineExprConstantFolder {
hasPoison_ = true;
return std::nullopt;
}
+ if (divideSignedWouldOverflow(lhs, rhs))
+ return std::nullopt;
return divideFloorSigned(lhs, rhs);
});
case AffineExprKind::CeilDiv:
@@ -81,6 +94,8 @@ class AffineExprConstantFolder {
hasPoison_ = true;
return std::nullopt;
}
+ if (divideSignedWouldOverflow(lhs, rhs))
+ return std::nullopt;
return divideCeilSigned(lhs, rhs);
});
case AffineExprKind::Constant:
diff --git a/mlir/test/Dialect/Affine/canonicalize.mlir b/mlir/test/Dialect/Affine/canonicalize.mlir
index 57b92e79403fe..d3f143c5fbacd 100644
--- a/mlir/test/Dialect/Affine/canonicalize.mlir
+++ b/mlir/test/Dialect/Affine/canonicalize.mlir
@@ -2653,3 +2653,60 @@ func.func @compose_into_access_keeps_alignment(%memref: memref<100xi32>, %i: ind
affine.store %val, %memref[%idx] { alignment = 16 } : memref<100xi32>
return
}
+
+// -----
+
+// Regression tests for https://github.com/llvm/llvm-project/issues/229706:
+// constant folding must not fold affine expressions whose result would
+// overflow int64_t.
+
+// CHECK: #[[$MAP:.*]] = affine_map<() -> (9223372036854775807 + 1)>
+// CHECK-LABEL: func @for_bound_add_overflow
+// CHECK: affine.for %{{.*}} = 0 to #[[$MAP]]() {
+func.func @for_bound_add_overflow(%m: memref<32xf32>) {
+ %max = arith.constant 9223372036854775807 : index
+ %cst = arith.constant 1.0 : f32
+ affine.for %j = 0 to affine_map<(d0) -> (d0 + 1)>(%max) {
+ affine.store %cst, %m[%j] : memref<32xf32>
+ }
+ return
+}
+
+// -----
+
+// CHECK: #[[$MAP:.*]] = affine_map<() -> (9223372036854775807 + 1)>
+// CHECK-LABEL: func @apply_add_overflow
+// CHECK: %[[RES:.+]] = affine.apply #[[$MAP]]()
+// CHECK: return %[[RES]]
+func.func @apply_add_overflow() -> index {
+ %max = arith.constant 9223372036854775807 : index
+ %0 = affine.apply affine_map<(d0) -> (d0 + 1)>(%max)
+ return %0 : index
+}
+
+// -----
+
+// CHECK: #[[$MAP:.*]] = affine_map<() -> (9223372036854775807 * 2)>
+// CHECK-LABEL: func @apply_mul_overflow
+// CHECK: %[[RES:.+]] = affine.apply #[[$MAP]]()
+// CHECK: return %[[RES]]
+func.func @apply_mul_overflow() -> index {
+ %max = arith.constant 9223372036854775807 : index
+ %0 = affine.apply affine_map<(d0) -> (d0 * 2)>(%max)
+ return %0 : index
+}
+
+// -----
+
+// CHECK-DAG: #[[$FLOOR:.*]] = affine_map<() -> (-9223372036854775808 floordiv -1)>
+// CHECK-DAG: #[[$CEIL:.*]] = affine_map<() -> (-9223372036854775808 ceildiv -1)>
+// CHECK-LABEL: func @apply_div_overflow
+// CHECK: %[[FLOOR_RES:.+]] = affine.apply #[[$FLOOR]]()
+// CHECK: %[[CEIL_RES:.+]] = affine.apply #[[$CEIL]]()
+// CHECK: return %[[FLOOR_RES]], %[[CEIL_RES]]
+func.func @apply_div_overflow() -> (index, index) {
+ %min = arith.constant -9223372036854775808 : index
+ %0 = affine.apply affine_map<(d0) -> (d0 floordiv -1)>(%min)
+ %1 = affine.apply affine_map<(d0) -> (d0 ceildiv -1)>(%min)
+ return %0, %1 : index, index
+}
>From 4891fb284949b85fb5738744b4c8a47471fd584d Mon Sep 17 00:00:00 2001
From: Mahmoud <mahmoud at Mahmouds-MacBook-Pro.local>
Date: Tue, 13 Oct 2026 11:44:18 +0300
Subject: [PATCH 2/2] [mlir][affine] Resolve conflicts with #228010
---
mlir/lib/IR/AffineExpr.cpp | 15 ++-
mlir/lib/IR/AffineMap.cpp | 1 +
mlir/test/Dialect/Affine/canonicalize.mlir | 114 ++++++++++-----------
3 files changed, 68 insertions(+), 62 deletions(-)
diff --git a/mlir/lib/IR/AffineExpr.cpp b/mlir/lib/IR/AffineExpr.cpp
index 209ad23f980b2..dd9a1f588a842 100644
--- a/mlir/lib/IR/AffineExpr.cpp
+++ b/mlir/lib/IR/AffineExpr.cpp
@@ -1354,11 +1354,14 @@ LogicalResult SimpleAffineExprFlattener::visitMulExpr(AffineBinaryOpExpr expr) {
return addLocalVariableSemiAffine(mulLhs, rhs, a * b, lhs, lhs.size());
}
- // Get the RHS constant. Bail out if any coefficient overflows.
+ // Get the RHS constant.
int64_t rhsConst = rhs[getConstantIndex()];
- for (int64_t &lhsElt : lhs)
- if (llvm::MulOverflow(lhsElt, rhsConst, lhsElt))
+ for (int64_t &lhsElt : lhs) {
+ auto [prod, overflow] = llvm::MulOverflow(lhsElt, rhsConst);
+ if (overflow)
return failure();
+ lhsElt = prod;
+ }
return success();
}
@@ -1368,10 +1371,12 @@ LogicalResult SimpleAffineExprFlattener::visitAddExpr(AffineBinaryOpExpr expr) {
const auto &rhs = operandExprStack.back();
auto &lhs = operandExprStack[operandExprStack.size() - 2];
assert(lhs.size() == rhs.size());
- // Update the LHS in place. Bail out if any coefficient overflows.
+ // Update the LHS in place.
for (unsigned i = 0, e = rhs.size(); i < e; i++) {
- if (llvm::AddOverflow(lhs[i], rhs[i], lhs[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/lib/IR/AffineMap.cpp b/mlir/lib/IR/AffineMap.cpp
index dfbf43392e8af..938a030fbf996 100644
--- a/mlir/lib/IR/AffineMap.cpp
+++ b/mlir/lib/IR/AffineMap.cpp
@@ -28,6 +28,7 @@ using llvm::divideSignedWouldOverflow;
using llvm::mod;
namespace {
+
// AffineExprConstantFolder evaluates an affine expression using constant
// operands passed in 'operandConsts'. Returns an IntegerAttr attribute
// representing the constant value of the affine expression evaluated on
diff --git a/mlir/test/Dialect/Affine/canonicalize.mlir b/mlir/test/Dialect/Affine/canonicalize.mlir
index d3f143c5fbacd..2724c2b0f009a 100644
--- a/mlir/test/Dialect/Affine/canonicalize.mlir
+++ b/mlir/test/Dialect/Affine/canonicalize.mlir
@@ -461,6 +461,63 @@ func.func @constant_fold_bounds(%N : index) {
// -----
+// Regression tests for https://github.com/llvm/llvm-project/issues/229706:
+// constant folding must not fold affine expressions whose result would
+// overflow int64_t.
+
+// CHECK: #[[$MAP:.*]] = affine_map<() -> (9223372036854775807 + 1)>
+// CHECK-LABEL: func @for_bound_add_overflow
+// CHECK: affine.for %{{.*}} = 0 to #[[$MAP]]() {
+func.func @for_bound_add_overflow(%m: memref<32xf32>) {
+ %max = arith.constant 9223372036854775807 : index
+ %cst = arith.constant 1.0 : f32
+ affine.for %j = 0 to affine_map<(d0) -> (d0 + 1)>(%max) {
+ affine.store %cst, %m[%j] : memref<32xf32>
+ }
+ return
+}
+
+// -----
+
+// CHECK: #[[$MAP:.*]] = affine_map<() -> (9223372036854775807 + 1)>
+// CHECK-LABEL: func @apply_add_overflow
+// CHECK: %[[RES:.+]] = affine.apply #[[$MAP]]()
+// CHECK: return %[[RES]]
+func.func @apply_add_overflow() -> index {
+ %max = arith.constant 9223372036854775807 : index
+ %0 = affine.apply affine_map<(d0) -> (d0 + 1)>(%max)
+ return %0 : index
+}
+
+// -----
+
+// CHECK: #[[$MAP:.*]] = affine_map<() -> (9223372036854775807 * 2)>
+// CHECK-LABEL: func @apply_mul_overflow
+// CHECK: %[[RES:.+]] = affine.apply #[[$MAP]]()
+// CHECK: return %[[RES]]
+func.func @apply_mul_overflow() -> index {
+ %max = arith.constant 9223372036854775807 : index
+ %0 = affine.apply affine_map<(d0) -> (d0 * 2)>(%max)
+ return %0 : index
+}
+
+// -----
+
+// CHECK-DAG: #[[$FLOOR:.*]] = affine_map<() -> (-9223372036854775808 floordiv -1)>
+// CHECK-DAG: #[[$CEIL:.*]] = affine_map<() -> (-9223372036854775808 ceildiv -1)>
+// CHECK-LABEL: func @apply_div_overflow
+// CHECK: %[[FLOOR_RES:.+]] = affine.apply #[[$FLOOR]]()
+// CHECK: %[[CEIL_RES:.+]] = affine.apply #[[$CEIL]]()
+// CHECK: return %[[FLOOR_RES]], %[[CEIL_RES]]
+func.func @apply_div_overflow() -> (index, index) {
+ %min = arith.constant -9223372036854775808 : index
+ %0 = affine.apply affine_map<(d0) -> (d0 floordiv -1)>(%min)
+ %1 = affine.apply affine_map<(d0) -> (d0 ceildiv -1)>(%min)
+ return %0, %1 : index, index
+}
+
+// -----
+
// CHECK-LABEL: func @fold_empty_loops()
func.func @fold_empty_loops() -> index {
%c0 = arith.constant 0 : index
@@ -2653,60 +2710,3 @@ func.func @compose_into_access_keeps_alignment(%memref: memref<100xi32>, %i: ind
affine.store %val, %memref[%idx] { alignment = 16 } : memref<100xi32>
return
}
-
-// -----
-
-// Regression tests for https://github.com/llvm/llvm-project/issues/229706:
-// constant folding must not fold affine expressions whose result would
-// overflow int64_t.
-
-// CHECK: #[[$MAP:.*]] = affine_map<() -> (9223372036854775807 + 1)>
-// CHECK-LABEL: func @for_bound_add_overflow
-// CHECK: affine.for %{{.*}} = 0 to #[[$MAP]]() {
-func.func @for_bound_add_overflow(%m: memref<32xf32>) {
- %max = arith.constant 9223372036854775807 : index
- %cst = arith.constant 1.0 : f32
- affine.for %j = 0 to affine_map<(d0) -> (d0 + 1)>(%max) {
- affine.store %cst, %m[%j] : memref<32xf32>
- }
- return
-}
-
-// -----
-
-// CHECK: #[[$MAP:.*]] = affine_map<() -> (9223372036854775807 + 1)>
-// CHECK-LABEL: func @apply_add_overflow
-// CHECK: %[[RES:.+]] = affine.apply #[[$MAP]]()
-// CHECK: return %[[RES]]
-func.func @apply_add_overflow() -> index {
- %max = arith.constant 9223372036854775807 : index
- %0 = affine.apply affine_map<(d0) -> (d0 + 1)>(%max)
- return %0 : index
-}
-
-// -----
-
-// CHECK: #[[$MAP:.*]] = affine_map<() -> (9223372036854775807 * 2)>
-// CHECK-LABEL: func @apply_mul_overflow
-// CHECK: %[[RES:.+]] = affine.apply #[[$MAP]]()
-// CHECK: return %[[RES]]
-func.func @apply_mul_overflow() -> index {
- %max = arith.constant 9223372036854775807 : index
- %0 = affine.apply affine_map<(d0) -> (d0 * 2)>(%max)
- return %0 : index
-}
-
-// -----
-
-// CHECK-DAG: #[[$FLOOR:.*]] = affine_map<() -> (-9223372036854775808 floordiv -1)>
-// CHECK-DAG: #[[$CEIL:.*]] = affine_map<() -> (-9223372036854775808 ceildiv -1)>
-// CHECK-LABEL: func @apply_div_overflow
-// CHECK: %[[FLOOR_RES:.+]] = affine.apply #[[$FLOOR]]()
-// CHECK: %[[CEIL_RES:.+]] = affine.apply #[[$CEIL]]()
-// CHECK: return %[[FLOOR_RES]], %[[CEIL_RES]]
-func.func @apply_div_overflow() -> (index, index) {
- %min = arith.constant -9223372036854775808 : index
- %0 = affine.apply affine_map<(d0) -> (d0 floordiv -1)>(%min)
- %1 = affine.apply affine_map<(d0) -> (d0 ceildiv -1)>(%min)
- return %0, %1 : index, index
-}
More information about the Mlir-commits
mailing list