[Mlir-commits] [mlir] [mlir][arith] Fix overflow bug in arith::CeilDivSIOp::fold (PR #90947)
Andrzej Warzyński
llvmlistbot at llvm.org
Mon May 6 06:58:10 PDT 2024
https://github.com/banach-space updated https://github.com/llvm/llvm-project/pull/90947
>From 9460305d2cfbc179563d994275a194f3da6be116 Mon Sep 17 00:00:00 2001
From: Andrzej Warzynski <andrzej.warzynski at arm.com>
Date: Fri, 3 May 2024 07:42:57 +0000
Subject: [PATCH 1/5] [mlir][arith] Fix overflow bug in
arith::CeilDivSIOp::fold
The folder for arith::CeilDivSIOp should only be applied when it can be
guaranteed that no overflow would happen. The current implementation
works fine when both dividends are positive and the only arithmetic
operation is the division itself.
However, in cases where at least one of the dividends is negative, the
division is split into multiple operations, e.g.: `- ( -a / b)`. That's
additional 2 operations on top of the actual division that can overflow
- the folder should check all 3 ops for overflow. The current logic
doesn't do that - it effectively only the last operation (i.e. the
division). It breaks when using e.g. MININT values (e.g. -128 for
8-bit integers) - negating such values overflows.
This PR makes sure that no folding happens if any of the intermediate
arithmetic operations overflows.
---
mlir/lib/Dialect/Arith/IR/ArithOps.cpp | 30 +++++++++++++++++--------
mlir/test/Transforms/constant-fold.mlir | 26 +++++++++++++++++++++
2 files changed, 47 insertions(+), 9 deletions(-)
diff --git a/mlir/lib/Dialect/Arith/IR/ArithOps.cpp b/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
index 6f995b93bc3ecd..a846f273996134 100644
--- a/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
+++ b/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
@@ -701,22 +701,34 @@ OpFoldResult arith::CeilDivSIOp::fold(FoldAdaptor adaptor) {
// Both positive, return ceil(a, b).
return signedCeilNonnegInputs(a, b, overflowOrDiv0);
}
+
+ bool overflowNegA = false;
+ bool overflowNegB = false;
+ bool overflowNegDiv = false;
+ bool overflowDiv = false;
if (!aGtZero && !bGtZero) {
// Both negative, return ceil(-a, -b).
- APInt posA = zero.ssub_ov(a, overflowOrDiv0);
- APInt posB = zero.ssub_ov(b, overflowOrDiv0);
- return signedCeilNonnegInputs(posA, posB, overflowOrDiv0);
+ APInt posA = zero.ssub_ov(a, overflowNegA);
+ APInt posB = zero.ssub_ov(b, overflowNegB);
+ APInt res = signedCeilNonnegInputs(posA, posB, overflowDiv);
+ overflowOrDiv0 = (overflowNegA || overflowNegB || overflowDiv);
+ return res;
}
if (!aGtZero && bGtZero) {
// A is negative, b is positive, return - ( -a / b).
- APInt posA = zero.ssub_ov(a, overflowOrDiv0);
- APInt div = posA.sdiv_ov(b, overflowOrDiv0);
- return zero.ssub_ov(div, overflowOrDiv0);
+ APInt posA = zero.ssub_ov(a, overflowNegA);
+ APInt div = posA.sdiv_ov(b, overflowDiv);
+ APInt res = zero.ssub_ov(div, overflowNegDiv);
+ overflowOrDiv0 = (overflowNegA || overflowDiv || overflowNegDiv);
+ return res;
}
// A is positive, b is negative, return - (a / -b).
- APInt posB = zero.ssub_ov(b, overflowOrDiv0);
- APInt div = a.sdiv_ov(posB, overflowOrDiv0);
- return zero.ssub_ov(div, overflowOrDiv0);
+ APInt posB = zero.ssub_ov(b, overflowNegB);
+ APInt div = a.sdiv_ov(posB, overflowDiv);
+ APInt res = zero.ssub_ov(div, overflowNegDiv);
+
+ overflowOrDiv0 = (overflowNegB || overflowDiv || overflowNegDiv);
+ return res;
});
return overflowOrDiv0 ? Attribute() : result;
diff --git a/mlir/test/Transforms/constant-fold.mlir b/mlir/test/Transforms/constant-fold.mlir
index 253163f2af9110..be0650a9a92043 100644
--- a/mlir/test/Transforms/constant-fold.mlir
+++ b/mlir/test/Transforms/constant-fold.mlir
@@ -478,6 +478,32 @@ func.func @simple_arith.ceildivsi() -> (i32, i32, i32, i32, i32) {
// -----
+// CHECK-LABEL: func @simple_arith.ceildivsi_overflow
+func.func @simple_arith.ceildivsi_overflow() -> (i8, i16, i32) {
+ // The negative values below are MININTs for the corresponding bit-width. The
+ // folder will try to negate them (so that the division operartes on two
+ // positive numbers), but that would cause overflow (negating MININT
+ // overflows). Hence folding should not happen and the original ceildivsi is
+ // preserved.
+
+ // CHECK-COUNT-3: arith.ceildivsi
+ %0 = arith.constant 7 : i8
+ %min_int_i8 = arith.constant -128 : i8
+ %2 = arith.ceildivsi %min_int_i8, %0 : i8
+
+ %3 = arith.constant 7 : i16
+ %min_int_i16 = arith.constant -32768 : i16
+ %5 = arith.ceildivsi %min_int_i16, %3 : i16
+
+ %6 = arith.constant 7 : i32
+ %min_int_i32 = arith.constant -2147483648 : i32
+ %8 = arith.ceildivsi %min_int_i32, %6 : i32
+
+ return %2, %5, %8 : i8, i16, i32
+}
+
+// -----
+
// CHECK-LABEL: func @simple_arith.ceildivui
func.func @simple_arith.ceildivui() -> (i32, i32, i32, i32, i32) {
// CHECK-DAG: [[C0:%.+]] = arith.constant 0
>From 12f7ed46b8de232f937daa77f1fb6e8f713787d1 Mon Sep 17 00:00:00 2001
From: Andrzej Warzynski <andrzej.warzynski at arm.com>
Date: Fri, 3 May 2024 11:24:55 +0000
Subject: [PATCH 2/5] fixup! [mlir][arith] Fix overflow bug in
arith::CeilDivSIOp::fold
Add comment, update var name (overflowNegDiv -> overflowNegRes)
---
mlir/lib/Dialect/Arith/IR/ArithOps.cpp | 12 +++++++-----
1 file changed, 7 insertions(+), 5 deletions(-)
diff --git a/mlir/lib/Dialect/Arith/IR/ArithOps.cpp b/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
index a846f273996134..c2b8895bc32400 100644
--- a/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
+++ b/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
@@ -702,10 +702,12 @@ OpFoldResult arith::CeilDivSIOp::fold(FoldAdaptor adaptor) {
return signedCeilNonnegInputs(a, b, overflowOrDiv0);
}
+ // No folding happens if any of the intermediate arithmetic operations
+ // overflows.
bool overflowNegA = false;
bool overflowNegB = false;
- bool overflowNegDiv = false;
bool overflowDiv = false;
+ bool overflowNegRes = false;
if (!aGtZero && !bGtZero) {
// Both negative, return ceil(-a, -b).
APInt posA = zero.ssub_ov(a, overflowNegA);
@@ -718,16 +720,16 @@ OpFoldResult arith::CeilDivSIOp::fold(FoldAdaptor adaptor) {
// A is negative, b is positive, return - ( -a / b).
APInt posA = zero.ssub_ov(a, overflowNegA);
APInt div = posA.sdiv_ov(b, overflowDiv);
- APInt res = zero.ssub_ov(div, overflowNegDiv);
- overflowOrDiv0 = (overflowNegA || overflowDiv || overflowNegDiv);
+ APInt res = zero.ssub_ov(div, overflowNegRes);
+ overflowOrDiv0 = (overflowNegA || overflowDiv || overflowNegRes);
return res;
}
// A is positive, b is negative, return - (a / -b).
APInt posB = zero.ssub_ov(b, overflowNegB);
APInt div = a.sdiv_ov(posB, overflowDiv);
- APInt res = zero.ssub_ov(div, overflowNegDiv);
+ APInt res = zero.ssub_ov(div, overflowNegRes);
- overflowOrDiv0 = (overflowNegB || overflowDiv || overflowNegDiv);
+ overflowOrDiv0 = (overflowNegB || overflowDiv || overflowNegRes);
return res;
});
>From 0715f2a25090a9fb84db47ed1a23cc2169d4ce3e Mon Sep 17 00:00:00 2001
From: Andrzej Warzynski <andrzej.warzynski at arm.com>
Date: Sun, 5 May 2024 14:23:17 +0000
Subject: [PATCH 3/5] fixup! fixup! [mlir][arith] Fix overflow bug in
arith::CeilDivSIOp::fold
Harden the test
---
mlir/test/Transforms/constant-fold.mlir | 11 ++++++++++-
1 file changed, 10 insertions(+), 1 deletion(-)
diff --git a/mlir/test/Transforms/constant-fold.mlir b/mlir/test/Transforms/constant-fold.mlir
index be0650a9a92043..90307dd4d9970d 100644
--- a/mlir/test/Transforms/constant-fold.mlir
+++ b/mlir/test/Transforms/constant-fold.mlir
@@ -486,15 +486,24 @@ func.func @simple_arith.ceildivsi_overflow() -> (i8, i16, i32) {
// overflows). Hence folding should not happen and the original ceildivsi is
// preserved.
- // CHECK-COUNT-3: arith.ceildivsi
+ // CHECK-DAG: %[[C_1:.*]] = arith.constant 7 : i8
+ // CHECK-DAG: %[[MIN_I8:.*]] = arith.constant -128 : i8
+ // CHECK-DAG: %[[C_2:.*]] = arith.constant 7 : i16
+ // CHECK-DAG: %[[MIN_I16:.*]] = arith.constant -32768 : i16
+ // CHECK-DAG: %[[C_3:.*]] = arith.constant 7 : i32
+ // CHECK-DAG: %[[MIN_I32:.*]] = arith.constant -2147483648 : i32
+
+ // CHECK-NEXT: %[[CEILDIV_1:.*]] = arith.ceildivsi %[[MIN_I8]], %[[C_1]] : i8
%0 = arith.constant 7 : i8
%min_int_i8 = arith.constant -128 : i8
%2 = arith.ceildivsi %min_int_i8, %0 : i8
+ // CHECK-NEXT: %[[CEILDIV_2:.*]] = arith.ceildivsi %[[MIN_I16]], %[[C_2]] : i16
%3 = arith.constant 7 : i16
%min_int_i16 = arith.constant -32768 : i16
%5 = arith.ceildivsi %min_int_i16, %3 : i16
+ // CHECK-NEXT: %[[CEILDIV_2:.*]] = arith.ceildivsi %[[MIN_I32]], %[[C_3]] : i32
%6 = arith.constant 7 : i32
%min_int_i32 = arith.constant -2147483648 : i32
%8 = arith.ceildivsi %min_int_i32, %6 : i32
>From a08f4dae2966c0bfaeb65266673079e06e9938c0 Mon Sep 17 00:00:00 2001
From: Andrzej Warzynski <andrzej.warzynski at arm.com>
Date: Mon, 6 May 2024 08:57:50 +0000
Subject: [PATCH 4/5] fixup! fixup! fixup! [mlir][arith] Fix overflow bug in
arith::CeilDivSIOp::fold
Fix typo
---
mlir/test/Transforms/constant-fold.mlir | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/mlir/test/Transforms/constant-fold.mlir b/mlir/test/Transforms/constant-fold.mlir
index 90307dd4d9970d..d8c99be71cb461 100644
--- a/mlir/test/Transforms/constant-fold.mlir
+++ b/mlir/test/Transforms/constant-fold.mlir
@@ -481,7 +481,7 @@ func.func @simple_arith.ceildivsi() -> (i32, i32, i32, i32, i32) {
// CHECK-LABEL: func @simple_arith.ceildivsi_overflow
func.func @simple_arith.ceildivsi_overflow() -> (i8, i16, i32) {
// The negative values below are MININTs for the corresponding bit-width. The
- // folder will try to negate them (so that the division operartes on two
+ // folder will try to negate them (so that the division operates on two
// positive numbers), but that would cause overflow (negating MININT
// overflows). Hence folding should not happen and the original ceildivsi is
// preserved.
>From 260985778d1871198bdca0071e788712b07eec83 Mon Sep 17 00:00:00 2001
From: Andrzej Warzynski <andrzej.warzynski at arm.com>
Date: Mon, 6 May 2024 13:57:33 +0000
Subject: [PATCH 5/5] fixup! fixup! fixup! fixup! [mlir][arith] Fix overflow
bug in arith::CeilDivSIOp::fold
Add comments
---
mlir/lib/Dialect/Arith/IR/ArithOps.cpp | 2 ++
mlir/test/Transforms/constant-fold.mlir | 3 +++
2 files changed, 5 insertions(+)
diff --git a/mlir/lib/Dialect/Arith/IR/ArithOps.cpp b/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
index c2b8895bc32400..a1568d0ebba3a8 100644
--- a/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
+++ b/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
@@ -683,6 +683,8 @@ OpFoldResult arith::CeilDivSIOp::fold(FoldAdaptor adaptor) {
return getLhs();
// Don't fold if it would overflow or if it requires a division by zero.
+ // TODO: This hook won't fold operations where a = MININT, because
+ // negating MININT overflows. This can be improved.
bool overflowOrDiv0 = false;
auto result = constFoldBinaryOp<IntegerAttr>(
adaptor.getOperands(), [&](APInt a, const APInt &b) {
diff --git a/mlir/test/Transforms/constant-fold.mlir b/mlir/test/Transforms/constant-fold.mlir
index d8c99be71cb461..981757aed9b1d6 100644
--- a/mlir/test/Transforms/constant-fold.mlir
+++ b/mlir/test/Transforms/constant-fold.mlir
@@ -486,6 +486,9 @@ func.func @simple_arith.ceildivsi_overflow() -> (i8, i16, i32) {
// overflows). Hence folding should not happen and the original ceildivsi is
// preserved.
+ // TODO: The folder should be able to fold the following by avoiding
+ // intermediate operations that overflow.
+
// CHECK-DAG: %[[C_1:.*]] = arith.constant 7 : i8
// CHECK-DAG: %[[MIN_I8:.*]] = arith.constant -128 : i8
// CHECK-DAG: %[[C_2:.*]] = arith.constant 7 : i16
More information about the Mlir-commits
mailing list