[MLIR][Shape] Canonicalize subsequent `size_to_index` and `index_to_size`

Eliminate the subsequent applications of `size_to_index` and `index_to_size`.

Differential Revision: https://reviews.llvm.org/D82083

diff --git a/mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td b/mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td
index 21d76a3..1fb88a7 100644 (file)
--- a/mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td
+++ b/mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td
@@ -249,6 +249,7 @@ def Shape_IndexToSizeOp : Shape_Op<"index_to_size", [NoSideEffect]> {
   let assemblyFormat = "$arg attr-dict";
 
   let hasFolder = 1;
+  let hasCanonicalizer = 1;
 }
 
 def Shape_JoinOp : Shape_Op<"join", []> {

diff --git a/mlir/lib/Dialect/Shape/IR/Shape.cpp b/mlir/lib/Dialect/Shape/IR/Shape.cpp
index b1f8043..e251a38 100644 (file)
--- a/mlir/lib/Dialect/Shape/IR/Shape.cpp
+++ b/mlir/lib/Dialect/Shape/IR/Shape.cpp
@@ -397,6 +397,11 @@ OpFoldResult IndexToSizeOp::fold(ArrayRef<Attribute> operands) {
   return {};
 }
 
+void IndexToSizeOp::getCanonicalizationPatterns(
+    OwningRewritePatternList &patterns, MLIRContext *context) {
+  patterns.insert<SizeToIndexToSizeCanonicalization>(context);
+}
+
 //===----------------------------------------------------------------------===//
 // FromExtentsOp
 //===----------------------------------------------------------------------===//

diff --git a/mlir/lib/Dialect/Shape/IR/ShapeCanonicalization.td b/mlir/lib/Dialect/Shape/IR/ShapeCanonicalization.td
index 43ea27f..b03e512 100644 (file)
--- a/mlir/lib/Dialect/Shape/IR/ShapeCanonicalization.td
+++ b/mlir/lib/Dialect/Shape/IR/ShapeCanonicalization.td
@@ -22,3 +22,7 @@ def IndexToSizeToIndexCanonicalization : Pat<
   (Shape_SizeToIndexOp (Shape_IndexToSizeOp $arg)),
   (replaceWithValue $arg)>;
 
+def SizeToIndexToSizeCanonicalization : Pat<
+  (Shape_IndexToSizeOp (Shape_SizeToIndexOp $arg)),
+  (replaceWithValue $arg)>;
+

diff --git a/mlir/test/Dialect/Shape/canonicalize.mlir b/mlir/test/Dialect/Shape/canonicalize.mlir
index 1da1b70..1b9f392 100644 (file)
--- a/mlir/test/Dialect/Shape/canonicalize.mlir
+++ b/mlir/test/Dialect/Shape/canonicalize.mlir
@@ -503,3 +503,15 @@ func @index_to_size_to_index(%index : index) -> index {
   return %result : index
 }
 
+// -----
+
+// Canonicalize redundant conversion from `size` to `index` and back.
+// CHECK-LABEL: @size_to_index_to_size
+// CHECK-SAME: (%[[SIZE:.*]]: !shape.size) -> !shape.size
+func @size_to_index_to_size(%size : !shape.size) -> !shape.size {
+  // CHECK: return %[[SIZE]] : !shape.size
+  %idx = shape.size_to_index %size
+  %result = shape.index_to_size %idx
+  return %result : !shape.size
+}
+