[Mlir-commits] [mlir] bf2a4f3 - [MLIR][Shape] Canonicalize subsequent `index_to_size` and `size_to_index`

[Frederik Gossen]

Author: Frederik Gossen
Date: 2020-06-25T12:02:49Z
New Revision: bf2a4f3b3a0b31ce6cb5a031e696d3061e1e1160

URL: https://github.com/llvm/llvm-project/commit/bf2a4f3b3a0b31ce6cb5a031e696d3061e1e1160
DIFF: https://github.com/llvm/llvm-project/commit/bf2a4f3b3a0b31ce6cb5a031e696d3061e1e1160.diff

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

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

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

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td
    mlir/lib/Dialect/Shape/IR/Shape.cpp
    mlir/lib/Dialect/Shape/IR/ShapeCanonicalization.td
    mlir/test/Dialect/Shape/canonicalize.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td b/mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td
index 2430fe62f13b..21d76a3c14e7 100644
--- a/mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td
+++ b/mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td
@@ -392,6 +392,7 @@ def Shape_SizeToIndexOp : Shape_Op<"size_to_index", [NoSideEffect]> {
   let assemblyFormat = "$arg attr-dict";
 
   let hasFolder = 1;
+  let hasCanonicalizer = 1;
 }
 
 def Shape_YieldOp : Shape_Op<"yield",

diff  --git a/mlir/lib/Dialect/Shape/IR/Shape.cpp b/mlir/lib/Dialect/Shape/IR/Shape.cpp
index 2d952183050e..b1f804356f0e 100644
--- a/mlir/lib/Dialect/Shape/IR/Shape.cpp
+++ b/mlir/lib/Dialect/Shape/IR/Shape.cpp
@@ -536,6 +536,11 @@ OpFoldResult SizeToIndexOp::fold(ArrayRef<Attribute> operands) {
   return {};
 }
 
+void SizeToIndexOp::getCanonicalizationPatterns(
+    OwningRewritePatternList &patterns, MLIRContext *context) {
+  patterns.insert<IndexToSizeToIndexCanonicalization>(context);
+}
+
 //===----------------------------------------------------------------------===//
 // YieldOp
 //===----------------------------------------------------------------------===//

diff  --git a/mlir/lib/Dialect/Shape/IR/ShapeCanonicalization.td b/mlir/lib/Dialect/Shape/IR/ShapeCanonicalization.td
index 78c9119f1292..43ea27f169b5 100644
--- a/mlir/lib/Dialect/Shape/IR/ShapeCanonicalization.td
+++ b/mlir/lib/Dialect/Shape/IR/ShapeCanonicalization.td
@@ -9,6 +9,7 @@ def AllInputShapesEq : Constraint<CPred< [{
 }]>>;
 
 // Canonicalization patterns.
+
 def CstrBroadcastableEqOps : Pat<(Shape_CstrBroadcastableOp:$op $lhs, $rhs),
   (Shape_ConstWitnessOp ConstBoolAttrTrue),
   [(EqualBinaryOperands $lhs, $rhs)]>;
@@ -16,3 +17,8 @@ def CstrBroadcastableEqOps : Pat<(Shape_CstrBroadcastableOp:$op $lhs, $rhs),
 def CstrEqEqOps : Pat<(Shape_CstrEqOp:$op $shapes),
   (Shape_ConstWitnessOp ConstBoolAttrTrue),
   [(AllInputShapesEq $shapes)]>;
+
+def IndexToSizeToIndexCanonicalization : Pat<
+  (Shape_SizeToIndexOp (Shape_IndexToSizeOp $arg)),
+  (replaceWithValue $arg)>;
+

diff  --git a/mlir/test/Dialect/Shape/canonicalize.mlir b/mlir/test/Dialect/Shape/canonicalize.mlir
index 9fb48e6c896e..1da1b70a90a3 100644
--- a/mlir/test/Dialect/Shape/canonicalize.mlir
+++ b/mlir/test/Dialect/Shape/canonicalize.mlir
@@ -492,3 +492,14 @@ func @dont_canonicalize_rank(%arg : tensor<*xf32>) -> !shape.size {
 %rank = shape.rank %shape
 return %rank : !shape.size
 }
+
+// Canonicalize redundant conversion from `index` to `size` and back.
+// CHECK-LABEL: @index_to_size_to_index
+// CHECK-SAME: (%[[IDX:.*]]: index) -> index
+func @index_to_size_to_index(%index : index) -> index {
+  // CHECK: return %[[IDX]] : index
+  %size = shape.index_to_size %index
+  %result = shape.size_to_index %size
+  return %result : index
+}
+