From 66e0f66d8ffb9d889ccb3a411c2d1e04a5e0ec41 Mon Sep 17 00:00:00 2001
From: Frederik Gossen <frgossen@google.com>
Date: Thu, 25 Jun 2020 11:59:19 +0000
Subject: [PATCH] [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
---
 mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td     |  1 +
 mlir/lib/Dialect/Shape/IR/Shape.cpp                |  5 +++++
 mlir/lib/Dialect/Shape/IR/ShapeCanonicalization.td |  4 ++++
 mlir/test/Dialect/Shape/canonicalize.mlir          | 12 ++++++++++++
 4 files changed, 22 insertions(+)

diff --git a/mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td b/mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td
index 21d76a3..1fb88a7 100644
--- 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
--- 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
--- 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
--- 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
+}
+
-- 
2.7.4