[mlir] Add folding of memref_cast inside another memref_cast

[ftynse]

Summary

There exists a generic folding facility that folds the operand of a memref_cast
into users of memref_cast that support this. However, it was not used for the
memref_cast itself. Fix it to enable elimination of memref_cast chains such as

%1 = memref_cast %0 : A to B
%2 = memref_cast %1 : B to A

that is achieved by combining the folding with the existing "A to A" cast
elimination.

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

if (aT && bT) {
 
     return true;
   }
 
   return false;
 }
 
 OpFoldResult MemRefCastOp::fold(ArrayRef<Attribute> operands) {
-  return impl::foldCastOp(*this);
+  if (Value folded = impl::foldCastOp(*this))
+    return folded;
+  return succeeded(foldMemRefCast(*this)) ? getResult() : Value();
 }
 
 //===----------------------------------------------------------------------===//
 // MemRefReinterpretCastOp
 //===----------------------------------------------------------------------===//
 
 void mlir::MemRefReinterpretCastOp::build(
     OpBuilder &b, OperationState &result, MemRefType resultType, Value source,

mlir/test/Transforms/canonicalize.mlir

func @memref_cast_folding(%arg0: memref<4 x f32>, %arg1: f32) -> (f32, f32) {
 
   // CHECK-NEXT: dealloc %arg0 : memref<4xf32>
   dealloc %0: memref<?xf32>
 
   // CHECK-NEXT: return %{{.*}}
   return %1, %2 : f32, f32
 }
 
+// CHECK-LABEL: @fold_memref_cast_in_memref_cast
+// CHECK-SAME: (%[[ARG0:.*]]: memref<42x42xf64>)
+func @fold_memref_cast_in_memref_cast(%0: memref<42x42xf64>) {
+  // CHECK: %[[folded:.*]] = memref_cast %[[ARG0]] : memref<42x42xf64> to memref<?x?xf64>
+  %4 = memref_cast %0 : memref<42x42xf64> to memref<?x42xf64>
+  // CHECK-NOT: memref_cast
+  %5 = memref_cast %4 : memref<?x42xf64> to memref<?x?xf64>
+  // CHECK: "test.user"(%[[folded]])
+  "test.user"(%5) : (memref<?x?xf64>) -> ()
+  return
+}
+
+// CHECK-LABEL: @fold_memref_cast_chain
+// CHECK-SAME: (%[[ARG0:.*]]: memref<42x42xf64>)
+func @fold_memref_cast_chain(%0: memref<42x42xf64>) {
+  // CHECK-NOT: memref_cast
+  %4 = memref_cast %0 : memref<42x42xf64> to memref<?x42xf64>
+  %5 = memref_cast %4 : memref<?x42xf64> to memref<42x42xf64>
+  // CHECK: "test.user"(%[[ARG0]])
+  "test.user"(%5) : (memref<42x42xf64>) -> ()
+  return
+}
+
 // CHECK-LABEL: func @alloc_const_fold
 func @alloc_const_fold() -> memref<?xf32> {
   // CHECK-NEXT: %0 = alloc() : memref<4xf32>
   %c4 = constant 4 : index
   %a = alloc(%c4) : memref<?xf32>
 
   // CHECK-NEXT: %1 = memref_cast %0 : memref<4xf32> to memref<?xf32>
   // CHECK-NEXT: return %1 : memref<?xf32>