[mlir][arith] Fix overflow bug in arith::CeilDivSIOp::fold #90947
Conversation
|
@llvm/pr-subscribers-mlir-arith @llvm/pr-subscribers-mlir Author: Andrzej Warzyński (banach-space) ChangesThe folder for arith::CeilDivSIOp should only be applied when it can be However, in cases where at least one of the dividends is negative, the
This PR makes sure that no folding happens if any of the intermediate Full diff: https://github.com/llvm/llvm-project/pull/90947.diff 2 Files Affected:
diff --git a/mlir/lib/Dialect/Arith/IR/ArithOps.cpp b/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
index 6f995b93bc3ecd..a89634dfe07a2f 100644
--- a/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
+++ b/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
@@ -701,22 +701,35 @@ 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..8507342f59297f 100644
--- a/mlir/test/Transforms/constant-fold.mlir
+++ b/mlir/test/Transforms/constant-fold.mlir
@@ -478,6 +478,26 @@ 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) {
+ // CHECK-COUNT-3: arith.ceildivsi
+ %0 = arith.constant 7 : i8
+ %1 = arith.constant -128 : i8
+ %2 = arith.ceildivsi %1, %0 : i8
+
+ %3 = arith.constant 7 : i16
+ %4 = arith.constant -32768 : i16
+ %5 = arith.ceildivsi %4, %3 : i16
+
+ %6 = arith.constant 7 : i32
+ %7 = arith.constant -2147483648 : i32
+ %8 = arith.ceildivsi %7, %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
|
|
✅ With the latest revision this PR passed the C/C++ code formatter. |
banach-space
force-pushed
the
andrzej/fix_fold_overflow
branch
from
c6480bc to
cc36865
Compare
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.
banach-space
force-pushed
the
andrzej/fix_fold_overflow
branch
from
cc36865 to
9460305
Compare
Add comment, update var name (overflowNegDiv -> overflowNegRes)
|
Thanks for taking a look @Lewuathe 🙏🏻 |
There was a problem hiding this comment.
Do we need to update the wording around overflow conditions in the op docs? https://mlir.llvm.org/docs/Dialects/ArithOps/#arithceildivsi-arithceildivsiop this doesn't define what we understand by signed division overflow
|
I'm a little confused here (due to my own ignorance): the original bug was on the conversion pass Perhaps an extra integration test, consisting of an end-to-end run of the original file with |
Constant folding is always involved by the dialect conversion framework implicitly.
This shouldn't be needed: when we root cause the problem we focus on unit-testing instead. |
…SIOp::fold Fix typo
There was a problem hiding this comment.
This seems to me like it makes the folder strictly more conservative, which should be safe.
It is still open whether we couldn't be more aggressive though: are there cases where we could fold but we don't right now? (if so can you add a TODO? Possibly a test for these?)
Can we accurately catch the real overflow cases and folds to poison for example? (arith may not generate poison out-of-thin air for at the moment, which could be a blocker)
…CeilDivSIOp::fold Add comments
+1, but I'd rather defer to somebody with a stronger opinion and more experience with Arith.
Yes, we could be more aggressive. In fact, that's what @bviyer is proposing here: I suggested that first we make sure that we don't drop these overflow flags on the floor (i.e. that we land this PR first).
Yes, the cases in the test that I added (I've added TODOs). There might be more, but nothing obvious comes to mind. I don't see that many edge cases being tested for Arith - there might be even more room for improvement. |
|
I'll land this later today.
@kuhar, please let me know if you feel that this deserves a note in the docs. My intention was to refine/fix the implementation details rather than change the Op semantics. But I appreciate that in practice one defines the other (not sure whether that's intended though). |
Yeah I think adding a few examples of overflowing inputs there would go a long way and won't take much time to write. |
…arith::CeilDivSIOp::fold Add comments
Turns out this is already documented, so I only added a note documenting the changes made in this patch. @kuhar , wdyt? |
| signed division overflow. | ||
| Divison by zero, or signed division overflow (minimum value divided by -1) | ||
| is undefined behavior. While dividing minimum value by a value != -1 shouldn't | ||
| overflow, the current implementation treats it as such. When applied to `vector` |
There was a problem hiding this comment.
The sentence makes it look like dividing minimum value by a value != -1 is undefined behavior, I don't believe this is the intent.
I am not sure what you're trying to document here? Implementation details of the folder aren't relevant to the op semantics I believe.
|
@banach-space I wanted to clarify whether what we have there currently is the only kind of signed division overflow or if there are more (which I thought may be the case looking at the fold implementation);
If this is the only overflow kind, no need to update the docs. I haven't looked at the fold carefully, so apologies if I'm barking up the wrong tree, but the fold should handle all defined cases in the dialect doc. If it doesn't, we either need to make the semantics less defined or make the fold more robust. |
This is the TODO we discussed before I believe.
Why? It seems just fine to have "incomplete" or "conservative" folders: these are opportunistic, not guaranteed to fold all possible cases in general. |
Yes, I think we agree on this. That's what I meant by 'should' (vs. 'must'). And to reiterate, my only concern is if we are discovering more overflow cases that don't lower / fold properly yet are allowed by the op spec. If this is not the case, I don't have any issues with this patch. |
Agreed. It wasn't clear to me what the expectation was, so I proposed sth to progress the discussion. Like you said, it's an implementation detail, so let me revert that.
The cases that's being fixed here is not
I feel that we are all on the same page here, so I will land this later today. Please let me know if you feel differently. |
This reverts commit a9e8f3a (the previous commit)
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 either the dividend or divisor is negative (or both),
the division is split into multiple arith operations, e.g.:
- ( -a / b). That'sadditional 2 operations on top of the actual division that can overflow
The current logic doesn't do that - it effectively only checks 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.
Fixes #89382