[MLIR] Add folding constants canonicalization for mlir::index::AddOp. #111084
Conversation
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
Member
|
@llvm/pr-subscribers-mlir-index Author: weiwei chen (weiweichen) Changes
Full diff: https://github.com/llvm/llvm-project/pull/111084.diff 3 Files Affected:
diff --git a/mlir/include/mlir/Dialect/Index/IR/IndexOps.td b/mlir/include/mlir/Dialect/Index/IR/IndexOps.td
index a30ae9f739cbc6..ce1355316b09b8 100644
--- a/mlir/include/mlir/Dialect/Index/IR/IndexOps.td
+++ b/mlir/include/mlir/Dialect/Index/IR/IndexOps.td
@@ -56,6 +56,8 @@ def Index_AddOp : IndexBinaryOp<"add", [Commutative, Pure]> {
%c = index.add %a, %b
```
}];
+
+ let hasCanonicalizeMethod = 1;
}
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Dialect/Index/IR/IndexOps.cpp b/mlir/lib/Dialect/Index/IR/IndexOps.cpp
index 42401dae217ce1..ace9b43014a665 100644
--- a/mlir/lib/Dialect/Index/IR/IndexOps.cpp
+++ b/mlir/lib/Dialect/Index/IR/IndexOps.cpp
@@ -136,6 +136,45 @@ OpFoldResult AddOp::fold(FoldAdaptor adaptor) {
return {};
}
+/// Canonicalize
+/// ` x = v + c1; y = x + c2` to `x = v + (c1 + c2)`
+/// ` x = v + c1; y = c2 + x` to `x = v + (c1 + c2)`
+/// ` x = c1 + v; y = x + c2` to `x = v + (c1 + c2)`
+/// ` x = c1 + v; y = c2 + x` to `x = v + (c1 + c2)`
+LogicalResult AddOp::canonicalize(AddOp op, PatternRewriter &rewriter) {
+
+ auto matchConstant = [](mlir::index::AddOp op, Value &v, IntegerAttr &c) {
+ v = op.getLhs();
+ if (!mlir::matchPattern(op.getRhs(), mlir::m_Constant(&c))) {
+ v = op.getRhs();
+ if (!mlir::matchPattern(op.getLhs(), mlir::m_Constant(&c)))
+ return false;
+ }
+ return true;
+ };
+
+ IntegerAttr c1, c2;
+ Value v1, v2;
+
+ if (!matchConstant(op, v1, c1))
+ return rewriter.notifyMatchFailure(op.getLoc(),
+ "neither LHS nor RHS is constant");
+
+ auto add = v1.getDefiningOp<mlir::index::AddOp>();
+ if (!add)
+ return rewriter.notifyMatchFailure(op.getLoc(), "LHS is not a add");
+
+ if (!matchConstant(add, v2, c2))
+ return rewriter.notifyMatchFailure(op.getLoc(),
+ "neither LHS nor RHS is constant");
+
+ auto c = rewriter.create<mlir::index::ConstantOp>(op->getLoc(),
+ c1.getInt() + c2.getInt());
+ auto newAdd = rewriter.create<mlir::index::AddOp>(op->getLoc(), v2, c);
+
+ rewriter.replaceOp(op, newAdd);
+ return success();
+}
//===----------------------------------------------------------------------===//
// SubOp
diff --git a/mlir/test/Dialect/Index/index-canonicalize.mlir b/mlir/test/Dialect/Index/index-canonicalize.mlir
index 37aa33bfde952e..256e327e83ea9c 100644
--- a/mlir/test/Dialect/Index/index-canonicalize.mlir
+++ b/mlir/test/Dialect/Index/index-canonicalize.mlir
@@ -32,6 +32,21 @@ func.func @add_overflow() -> (index, index) {
return %2, %3 : index, index
}
+// CHECK-LABEL: @add
+func.func @add_fold_constants(%arg: index) -> (index) {
+ %0 = index.constant 1
+ %1 = index.constant 2
+ %2 = index.add %arg, %0
+ %3 = index.add %1, %2
+ %4 = index.add %3, %1
+ %5 = index.add %4, %0
+
+ // CHECK-DAG: [[A:%.*]] = index.constant 6
+ // CHECK-DAG: [[V0:%.*]] = index.add %arg0, [[A]]
+ // CHECK: return [[V0]]
+ return %5 : index
+}
+
// CHECK-LABEL: @sub
func.func @sub() -> index {
%0 = index.constant -2000000000
|
Member
|
@llvm/pr-subscribers-mlir Author: weiwei chen (weiweichen) Changes
Full diff: https://github.com/llvm/llvm-project/pull/111084.diff 3 Files Affected:
diff --git a/mlir/include/mlir/Dialect/Index/IR/IndexOps.td b/mlir/include/mlir/Dialect/Index/IR/IndexOps.td
index a30ae9f739cbc6..ce1355316b09b8 100644
--- a/mlir/include/mlir/Dialect/Index/IR/IndexOps.td
+++ b/mlir/include/mlir/Dialect/Index/IR/IndexOps.td
@@ -56,6 +56,8 @@ def Index_AddOp : IndexBinaryOp<"add", [Commutative, Pure]> {
%c = index.add %a, %b
```
}];
+
+ let hasCanonicalizeMethod = 1;
}
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Dialect/Index/IR/IndexOps.cpp b/mlir/lib/Dialect/Index/IR/IndexOps.cpp
index 42401dae217ce1..ace9b43014a665 100644
--- a/mlir/lib/Dialect/Index/IR/IndexOps.cpp
+++ b/mlir/lib/Dialect/Index/IR/IndexOps.cpp
@@ -136,6 +136,45 @@ OpFoldResult AddOp::fold(FoldAdaptor adaptor) {
return {};
}
+/// Canonicalize
+/// ` x = v + c1; y = x + c2` to `x = v + (c1 + c2)`
+/// ` x = v + c1; y = c2 + x` to `x = v + (c1 + c2)`
+/// ` x = c1 + v; y = x + c2` to `x = v + (c1 + c2)`
+/// ` x = c1 + v; y = c2 + x` to `x = v + (c1 + c2)`
+LogicalResult AddOp::canonicalize(AddOp op, PatternRewriter &rewriter) {
+
+ auto matchConstant = [](mlir::index::AddOp op, Value &v, IntegerAttr &c) {
+ v = op.getLhs();
+ if (!mlir::matchPattern(op.getRhs(), mlir::m_Constant(&c))) {
+ v = op.getRhs();
+ if (!mlir::matchPattern(op.getLhs(), mlir::m_Constant(&c)))
+ return false;
+ }
+ return true;
+ };
+
+ IntegerAttr c1, c2;
+ Value v1, v2;
+
+ if (!matchConstant(op, v1, c1))
+ return rewriter.notifyMatchFailure(op.getLoc(),
+ "neither LHS nor RHS is constant");
+
+ auto add = v1.getDefiningOp<mlir::index::AddOp>();
+ if (!add)
+ return rewriter.notifyMatchFailure(op.getLoc(), "LHS is not a add");
+
+ if (!matchConstant(add, v2, c2))
+ return rewriter.notifyMatchFailure(op.getLoc(),
+ "neither LHS nor RHS is constant");
+
+ auto c = rewriter.create<mlir::index::ConstantOp>(op->getLoc(),
+ c1.getInt() + c2.getInt());
+ auto newAdd = rewriter.create<mlir::index::AddOp>(op->getLoc(), v2, c);
+
+ rewriter.replaceOp(op, newAdd);
+ return success();
+}
//===----------------------------------------------------------------------===//
// SubOp
diff --git a/mlir/test/Dialect/Index/index-canonicalize.mlir b/mlir/test/Dialect/Index/index-canonicalize.mlir
index 37aa33bfde952e..256e327e83ea9c 100644
--- a/mlir/test/Dialect/Index/index-canonicalize.mlir
+++ b/mlir/test/Dialect/Index/index-canonicalize.mlir
@@ -32,6 +32,21 @@ func.func @add_overflow() -> (index, index) {
return %2, %3 : index, index
}
+// CHECK-LABEL: @add
+func.func @add_fold_constants(%arg: index) -> (index) {
+ %0 = index.constant 1
+ %1 = index.constant 2
+ %2 = index.add %arg, %0
+ %3 = index.add %1, %2
+ %4 = index.add %3, %1
+ %5 = index.add %4, %0
+
+ // CHECK-DAG: [[A:%.*]] = index.constant 6
+ // CHECK-DAG: [[V0:%.*]] = index.add %arg0, [[A]]
+ // CHECK: return [[V0]]
+ return %5 : index
+}
+
// CHECK-LABEL: @sub
func.func @sub() -> index {
%0 = index.constant -2000000000
|
Mogball
reviewed
Oct 8, 2024
| return {}; | ||
| } | ||
| /// Canonicalize | ||
| /// ` x = v + c1; y = x + c2` to `x = v + (c1 + c2)` |
Contributor
There was a problem hiding this comment.
This is the only permutation of the pattern that you actually need to implement. The canonicalizer will make sure that constants for commutative ops are always on the RHS.
Contributor
Author
There was a problem hiding this comment.
oh, good to know, PR updated!
…eic/canonicalize-index-add
Contributor
Author
|
@Mogball another 👀 please 🙏 ? |
Mogball
approved these changes
Oct 21, 2024
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
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.
mlir::index::AddOp.