[Mlir-commits] [mlir] 12df427 - [MLIR][Standard] Add folding for indexCast(indexCast(x)) -> x

[Stephen Neuendorffer]

Author: Stephen Neuendorffer
Date: 2020-02-10T10:23:56-08:00
New Revision: 12df427fb2225e1957771ba31e4e6e3350f2e4a9

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

LOG: [MLIR][Standard] Add folding for indexCast(indexCast(x)) -> x

Allow this only if the types are the same.  e.g.:
i16 -> index -> i16  or
index -> i16 -> index

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

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/StandardOps/Ops.td
    mlir/lib/Dialect/StandardOps/Ops.cpp
    mlir/test/Transforms/canonicalize.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/StandardOps/Ops.td b/mlir/include/mlir/Dialect/StandardOps/Ops.td
index a7647edd9909..0f4a1c7a0ea0 100644
--- a/mlir/include/mlir/Dialect/StandardOps/Ops.td
+++ b/mlir/include/mlir/Dialect/StandardOps/Ops.td
@@ -758,7 +758,7 @@ def IndexCastOp : CastOp<"index_cast">, Arguments<(ins AnyType:$in)> {
     static bool areCastCompatible(Type a, Type b);
   }];
 
-  let hasFolder = 0;
+  let hasFolder = 1;
 }
 
 def FPExtOp : CastOp<"fpext">, Arguments<(ins AnyType:$in)> {

diff  --git a/mlir/lib/Dialect/StandardOps/Ops.cpp b/mlir/lib/Dialect/StandardOps/Ops.cpp
index b1e0096c26e2..e7bfae7ed6db 100644
--- a/mlir/lib/Dialect/StandardOps/Ops.cpp
+++ b/mlir/lib/Dialect/StandardOps/Ops.cpp
@@ -1586,6 +1586,14 @@ bool IndexCastOp::areCastCompatible(Type a, Type b) {
          (a.isa<IntegerType>() && b.isIndex());
 }
 
+OpFoldResult IndexCastOp::fold(ArrayRef<Attribute> cstOperands) {
+  // Fold IndexCast(IndexCast(x)) -> x
+  auto cast = dyn_cast_or_null<IndexCastOp>(getOperand().getDefiningOp());
+  if (cast && cast.getOperand().getType() == getType())
+    return cast.getOperand();
+  return {};
+}
+
 //===----------------------------------------------------------------------===//
 // LoadOp
 //===----------------------------------------------------------------------===//

diff  --git a/mlir/test/Transforms/canonicalize.mlir b/mlir/test/Transforms/canonicalize.mlir
index 492e829a4639..fb6e9c433ae8 100644
--- a/mlir/test/Transforms/canonicalize.mlir
+++ b/mlir/test/Transforms/canonicalize.mlir
@@ -878,3 +878,12 @@ func @subview(%arg0 : index, %arg1 : index) -> (index, index) {
   // CHECK: return %[[C7]], %[[C11]]
   return %7, %8 : index, index
 }
+
+// CHECK-LABEL: func @index_cast
+// CHECK-SAME: %[[ARG_0:arg[0-9]+]]: i16
+func @index_cast(%arg0: i16) -> (i16) {
+  %11 = index_cast %arg0 : i16 to index
+  %12 = index_cast %11 : index to i16
+  // CHECK: return %[[ARG_0]] : i16
+  return %12 : i16
+}