[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

1637

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

mlir/include/mlir/Dialect/StandardOps/CMakeLists.txt

 set(LLVM_TARGET_DEFINITIONS Ops.td)
+mlir_tablegen(Ops.inc -gen-rewriters)
 mlir_tablegen(Ops.h.inc -gen-op-decls)
 mlir_tablegen(Ops.cpp.inc -gen-op-defs)
 mlir_tablegen(OpsEnums.h.inc -gen-enum-decls)
 mlir_tablegen(OpsEnums.cpp.inc -gen-enum-defs)
 add_public_tablegen_target(MLIRStandardOpsIncGen)

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 hasCanonicalizer = 1;
   let hasFolder = 0;
 }
+def TypesAreIdentical : Constraint<CPred<"$0.getType() == $1.getType()">>;
+def CombineIndexCastPattern :
+   Pat<(IndexCastOp:$out (IndexCastOp $in)),
+      (replaceWithValue $in),
+      [(TypesAreIdentical $in, $out)]>;
 
 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

 #include "mlir/Transforms/InliningUtils.h"
 #include "llvm/ADT/StringSwitch.h"
 #include "llvm/Support/FormatVariadic.h"
 #include "llvm/Support/raw_ostream.h"
 
 // Pull in all enum type definitions and utility function declarations.
 #include "mlir/Dialect/StandardOps/OpsEnums.cpp.inc"
 
+namespace mlir {
+#include "mlir/Dialect/StandardOps/Ops.inc"
+}
+
 using namespace mlir;
 
 //===----------------------------------------------------------------------===//
 // StandardOpsDialect Interfaces
 //===----------------------------------------------------------------------===//
 namespace {
 /// This class defines the interface for handling inlining with standard
 /// operations.

 //===----------------------------------------------------------------------===//
 
 // 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());
 }
 
+void IndexCastOp::getCanonicalizationPatterns(OwningRewritePatternList &results,
+                                              MLIRContext *context) {
+  /// indexCast(indexCast x) -> x
+  results.insert<CombineIndexCastPattern>(context);
+}
+
 //===----------------------------------------------------------------------===//
 // 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
+func @index_cast(%arg0: i16) -> (i16) {
+  %11 = index_cast %arg0 : i16 to index
+  %12 = index_cast %11 : index to i16
+  // CHECK: return %arg0 : i16
+  return %12 : i16
+}