[mlir][arith] Fold min/max with absorption and redundancy #160224
+353
−1
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -1116,6 +1116,18 @@ OpFoldResult arith::MaximumFOp::fold(FoldAdaptor adaptor) { | |
| if (matchPattern(adaptor.getRhs(), m_NegInfFloat())) | ||
| return getLhs(); | ||
|
|
||
| // max(max(a, b), b) -> max(a, b) | ||
| // max(max(a, b), a) -> max(a, b) | ||
| if (auto max = getLhs().getDefiningOp<MaximumFOp>()) | ||
| if (getRhs() == max.getRhs() || getRhs() == max.getLhs()) | ||
| return getLhs(); | ||
|
|
||
| // max(a, max(a, b)) -> max(a, b) | ||
| // max(b, max(a, b)) -> max(a, b) | ||
| if (auto max = getRhs().getDefiningOp<MaximumFOp>()) | ||
| if (getLhs() == max.getRhs() || getLhs() == max.getLhs()) | ||
| return getRhs(); | ||
|
|
||
| return constFoldBinaryOp<FloatAttr>( | ||
| adaptor.getOperands(), | ||
| [](const APFloat &a, const APFloat &b) { return llvm::maximum(a, b); }); | ||
|
|
@@ -1134,6 +1146,18 @@ OpFoldResult arith::MaxNumFOp::fold(FoldAdaptor adaptor) { | |
| if (matchPattern(adaptor.getRhs(), m_NaNFloat())) | ||
| return getLhs(); | ||
|
|
||
| // max(max(a, b), b) -> max(a, b) | ||
| // max(max(a, b), a) -> max(a, b) | ||
| if (auto max = getLhs().getDefiningOp<MaxNumFOp>()) | ||
| if (getRhs() == max.getRhs() || getRhs() == max.getLhs()) | ||
| return getLhs(); | ||
|
|
||
| // max(a, max(a, b)) -> max(a, b) | ||
| // max(b, max(a, b)) -> max(a, b) | ||
| if (auto max = getRhs().getDefiningOp<MaxNumFOp>()) | ||
| if (getLhs() == max.getRhs() || getLhs() == max.getLhs()) | ||
| return getRhs(); | ||
|
|
||
| return constFoldBinaryOp<FloatAttr>(adaptor.getOperands(), llvm::maxnum); | ||
| } | ||
|
|
||
|
|
@@ -1156,6 +1180,30 @@ OpFoldResult MaxSIOp::fold(FoldAdaptor adaptor) { | |
| return getLhs(); | ||
| } | ||
|
|
||
| // max(max(a, b), b) -> max(a, b) | ||
| // max(max(a, b), a) -> max(a, b) | ||
| if (auto max = getLhs().getDefiningOp<MaxSIOp>()) | ||
| if (getRhs() == max.getRhs() || getRhs() == max.getLhs()) | ||
| return getLhs(); | ||
|
|
||
| // max(a, max(a, b)) -> max(a, b) | ||
| // max(b, max(a, b)) -> max(a, b) | ||
|
Comment on lines
+1183
to
+1190
|
||
| if (auto max = getRhs().getDefiningOp<MaxSIOp>()) | ||
| if (getLhs() == max.getRhs() || getLhs() == max.getLhs()) | ||
| return getRhs(); | ||
|
|
||
| // max(min(a, b), a) -> a | ||
| // max(min(b, a), a) -> a | ||
| if (auto min = getLhs().getDefiningOp<MinSIOp>()) | ||
| if (getRhs() == min.getRhs() || getRhs() == min.getLhs()) | ||
| return getRhs(); | ||
|
|
||
| // max(a, min(a, b)) -> a | ||
| // max(a, min(b, a)) -> a | ||
|
Comment on lines
+1195
to
+1202
Member
There was a problem hiding this comment. I was wondering if this is correct when one of the values is |
||
| if (auto min = getRhs().getDefiningOp<MinSIOp>()) | ||
| if (getLhs() == min.getRhs() || getLhs() == min.getLhs()) | ||
| return getLhs(); | ||
|
|
||
| return constFoldBinaryOp<IntegerAttr>(adaptor.getOperands(), | ||
| [](const APInt &a, const APInt &b) { | ||
| return llvm::APIntOps::smax(a, b); | ||
|
|
@@ -1181,6 +1229,30 @@ OpFoldResult MaxUIOp::fold(FoldAdaptor adaptor) { | |
| return getLhs(); | ||
| } | ||
|
|
||
| // max(max(a, b), b) -> max(a, b) | ||
| // max(max(a, b), a) -> max(a, b) | ||
| if (auto max = getLhs().getDefiningOp<MaxUIOp>()) | ||
| if (getRhs() == max.getRhs() || getRhs() == max.getLhs()) | ||
| return getLhs(); | ||
|
|
||
| // max(a, max(a, b)) -> max(a, b) | ||
| // max(b, max(a, b)) -> max(a, b) | ||
| if (auto max = getRhs().getDefiningOp<MaxUIOp>()) | ||
| if (getLhs() == max.getRhs() || getLhs() == max.getLhs()) | ||
| return getRhs(); | ||
|
|
||
| // max(min(a, b), a) -> a | ||
| // max(min(b, a), a) -> a | ||
| if (auto min = getLhs().getDefiningOp<MinUIOp>()) | ||
| if (getRhs() == min.getRhs() || getRhs() == min.getLhs()) | ||
| return getRhs(); | ||
|
|
||
| // max(a, min(a, b)) -> a | ||
| // max(a, min(b, a)) -> a | ||
| if (auto min = getRhs().getDefiningOp<MinUIOp>()) | ||
| if (getLhs() == min.getRhs() || getLhs() == min.getLhs()) | ||
| return getLhs(); | ||
|
|
||
| return constFoldBinaryOp<IntegerAttr>(adaptor.getOperands(), | ||
| [](const APInt &a, const APInt &b) { | ||
| return llvm::APIntOps::umax(a, b); | ||
|
|
@@ -1200,6 +1272,18 @@ OpFoldResult arith::MinimumFOp::fold(FoldAdaptor adaptor) { | |
| if (matchPattern(adaptor.getRhs(), m_PosInfFloat())) | ||
| return getLhs(); | ||
|
|
||
| // min(min(a, b), b) -> min(a, b) | ||
| // min(min(a, b), a) -> min(a, b) | ||
| if (auto min = getLhs().getDefiningOp<MinimumFOp>()) | ||
| if (getRhs() == min.getRhs() || getRhs() == min.getLhs()) | ||
| return getLhs(); | ||
|
|
||
| // min(a, min(a, b)) -> min(a, b) | ||
| // min(b, min(a, b)) -> min(a, b) | ||
| if (auto min = getRhs().getDefiningOp<MinimumFOp>()) | ||
| if (getLhs() == min.getRhs() || getLhs() == min.getLhs()) | ||
| return getRhs(); | ||
|
|
||
| return constFoldBinaryOp<FloatAttr>( | ||
| adaptor.getOperands(), | ||
| [](const APFloat &a, const APFloat &b) { return llvm::minimum(a, b); }); | ||
|
|
@@ -1218,6 +1302,18 @@ OpFoldResult arith::MinNumFOp::fold(FoldAdaptor adaptor) { | |
| if (matchPattern(adaptor.getRhs(), m_NaNFloat())) | ||
| return getLhs(); | ||
|
|
||
| // min(min(a, b), b) -> min(a, b) | ||
| // min(min(a, b), a) -> min(a, b) | ||
| if (auto min = getLhs().getDefiningOp<MinNumFOp>()) | ||
| if (getRhs() == min.getRhs() || getRhs() == min.getLhs()) | ||
| return getLhs(); | ||
|
|
||
| // min(a, min(a, b)) -> min(a, b) | ||
| // min(b, min(a, b)) -> min(a, b) | ||
| if (auto min = getRhs().getDefiningOp<MinNumFOp>()) | ||
| if (getLhs() == min.getRhs() || getLhs() == min.getLhs()) | ||
| return getRhs(); | ||
|
|
||
| return constFoldBinaryOp<FloatAttr>( | ||
| adaptor.getOperands(), | ||
| [](const APFloat &a, const APFloat &b) { return llvm::minnum(a, b); }); | ||
|
|
@@ -1242,6 +1338,30 @@ OpFoldResult MinSIOp::fold(FoldAdaptor adaptor) { | |
| return getLhs(); | ||
| } | ||
|
|
||
| // min(min(a, b), b) -> min(a, b) | ||
| // min(min(a, b), a) -> min(a, b) | ||
| if (auto min = getLhs().getDefiningOp<MinSIOp>()) | ||
| if (getRhs() == min.getRhs() || getRhs() == min.getLhs()) | ||
| return getLhs(); | ||
|
|
||
| // min(a, min(a, b)) -> min(a, b) | ||
| // min(b, min(a, b)) -> min(a, b) | ||
| if (auto min = getRhs().getDefiningOp<MinSIOp>()) | ||
| if (getLhs() == min.getRhs() || getLhs() == min.getLhs()) | ||
| return getRhs(); | ||
|
|
||
| // min(max(a, b), a) -> a | ||
| // min(max(b, a), a) -> a | ||
| if (auto max = getLhs().getDefiningOp<MaxSIOp>()) | ||
| if (getRhs() == max.getRhs() || getRhs() == max.getLhs()) | ||
| return getRhs(); | ||
|
|
||
| // min(a, max(a, b)) -> a | ||
| // min(a, max(b, a)) -> a | ||
| if (auto max = getRhs().getDefiningOp<MaxSIOp>()) | ||
| if (getLhs() == max.getRhs() || getLhs() == max.getLhs()) | ||
| return getLhs(); | ||
|
|
||
| return constFoldBinaryOp<IntegerAttr>(adaptor.getOperands(), | ||
| [](const APInt &a, const APInt &b) { | ||
| return llvm::APIntOps::smin(a, b); | ||
|
|
@@ -1267,6 +1387,30 @@ OpFoldResult MinUIOp::fold(FoldAdaptor adaptor) { | |
| return getLhs(); | ||
| } | ||
|
|
||
| // min(min(a, b), b) -> min(a, b) | ||
| // min(min(a, b), a) -> min(a, b) | ||
| if (auto min = getLhs().getDefiningOp<MinUIOp>()) | ||
| if (getRhs() == min.getRhs() || getRhs() == min.getLhs()) | ||
| return getLhs(); | ||
|
|
||
| // min(a, min(a, b)) -> min(a, b) | ||
| // min(b, min(a, b)) -> min(a, b) | ||
| if (auto min = getRhs().getDefiningOp<MinUIOp>()) | ||
| if (getLhs() == min.getRhs() || getLhs() == min.getLhs()) | ||
| return getRhs(); | ||
|
|
||
| // min(max(a, b), a) -> a | ||
| // min(max(b, a), a) -> a | ||
| if (auto max = getLhs().getDefiningOp<MaxUIOp>()) | ||
| if (getRhs() == max.getRhs() || getRhs() == max.getLhs()) | ||
| return getRhs(); | ||
|
|
||
| // min(a, max(a, b)) -> a | ||
| // min(a, max(b, a)) -> a | ||
| if (auto max = getRhs().getDefiningOp<MaxUIOp>()) | ||
| if (getLhs() == max.getRhs() || getLhs() == max.getLhs()) | ||
| return getLhs(); | ||
|
|
||
| return constFoldBinaryOp<IntegerAttr>(adaptor.getOperands(), | ||
| [](const APInt &a, const APInt &b) { | ||
| return llvm::APIntOps::umin(a, b); | ||
|
|
||
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I wonder if instead of covering both cases here and for other ops, we should instead use the fact that these are commutative and first move the other
min/maxoperand to the RHS?Uh oh!
There was an error while loading. Please reload this page.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Good point. Currently, the Commutative trait only provides a generic fold that moves constant operands to the end:
llvm-project/mlir/lib/IR/Operation.cpp
Lines 856 to 873 in cac0635
so it helps in canonicalize patterns like
max(max(x, c0), c1) -> max(x, max(c0, c1)), where there are constants.But for cases like fold patterns in this pr
max(max(a, b), a) -> max(a, b), no constants are involved, so the commutative fold doesn’t normalize operands. That’s why I explicitly check both lhs and rhs here.An alternative would be to normalize operands inside the fold itself (e.g., always move the nested max to the RHS), but for now I kept both conditions for clarity.
cc @joker-eph
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Yes, I suggested this in the other revision: doing this in the trait would align all the commutative operation and then in-turn simplify all the folders (to avoid checking redundant forms)