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

[stephenneuendorffer]

Summary

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

[merge_guards_bot]

Unit tests: pass. 62315 tests passed, 0 failed and 838 were skipped.

clang-tidy: fail. clang-tidy found 0 errors and 1 warnings. 0 of them are added as review comments below (why?).

clang-format: fail. Please format your changes with clang-format by running git-clang-format HEAD^ or applying this patch.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[rriddle]

Can you implement this as a fold instead? 'fold' is usable in many more places than canonicalization patterns, e.g. DialectConversion.

[merge_guards_bot]

Unit tests: pass. 62315 tests passed, 0 failed and 838 were skipped.

clang-tidy: fail. clang-tidy found 0 errors and 1 warnings. 0 of them are added as review comments below (why?).

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[ftynse]

Looks good in general, please follow @rriddle's recommendation.

[nicolasvasilache]

+1 to what River said.

QQ, I could also see i16 -> index -> i32 => i16 -> i32.
Would it make sense to put here or do we prefer a separate revision?

[stephenneuendorffer]

In D73671#1852510, @nicolasvasilache wrote:

+1 to what River said.

QQ, I could also see i16 -> index -> i32 => i16 -> i32.
Would it make sense to put here or do we prefer a separate revision?

That cast is currently not legal as an IndexCastOp.

[stephenneuendorffer]

In D73671#1848442, @rriddle wrote:

Can you implement this as a fold instead? 'fold' is usable in many more places than canonicalization patterns, e.g. DialectConversion.

Any chance you can write something that would clarify this here: https://mlir.llvm.org/docs/Canonicalization/?

[merge_guards_bot]

Unit tests: pass. 62518 tests passed, 0 failed and 844 were skipped.

clang-tidy: unknown.

clang-format: unknown.

Build artifacts: CMakeCache.txt, console-log.txt, test-results.xml

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[rriddle]

mlir/lib/Dialect/StandardOps/Ops.cpp

1592

nit: Remove trivial braces and the extra newline at the end of this function.

mlir/include/mlir/Dialect/StandardOps/Ops.td

def IndexCastOp : CastOp<"index_cast">, Arguments<(ins AnyType:$in)> {
   }];
 
   let extraClassDeclaration = [{
     /// Return true if `a` and `b` are valid operand and result pairs for
     /// the operation.
     static bool areCastCompatible(Type a, Type b);
   }];
 
-  let hasFolder = 0;
+  let hasFolder = 1;
 }
 
 def FPExtOp : CastOp<"fpext">, Arguments<(ins AnyType:$in)> {
   let summary = "cast from floating-point to wider floating-point";
   let description = [{
     Cast a floating-point value to a larger floating-point-typed value.
     The destination type must to be strictly wider than the source type.
     Only scalars are currently supported.

mlir/lib/Dialect/StandardOps/Ops.cpp

 //===----------------------------------------------------------------------===//
 
 // Index cast is applicable from index to integer and backwards.
 bool IndexCastOp::areCastCompatible(Type a, Type b) {
   return (a.isIndex() && b.isa<IntegerType>()) ||
          (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
 //===----------------------------------------------------------------------===//
 
 static void print(OpAsmPrinter &p, LoadOp op) {
   p << "load " << op.getMemRef() << '[' << op.getIndices() << ']';
   p.printOptionalAttrDict(op.getAttrs());
   p << " : " << op.getMemRefType();

mlir/test/Transforms/canonicalize.mlir

func @subview(%arg0 : index, %arg1 : index) -> (index, index) {
   %7 = dim %4, 0 : memref<?x?x?xf32,
        affine_map<(d0, d1, d2)[s0, s1, s2, s3] -> (d0 * s1 + d1 * s2 + d2 * s3 + s0)>>
   %8 = dim %4, 1 : memref<?x?x?xf32,
        affine_map<(d0, d1, d2)[s0, s1, s2, s3] -> (d0 * s1 + d1 * s2 + d2 * s3 + s0)>>
 
   // 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
+}