[mlir][memref] Fix expanded shape ops memref.cast folding with changed type

[bkramer]

Summary

memref.expand_shape has verification logic to make sure
result dim must be static if all the collapsing src dims are static.

This can be relaxed once expand_shape supports more dynamism.

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

 
 void CollapseShapeOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                                   MLIRContext *context) {
   results.add<CollapseReshapeOps<CollapseShapeOp>,
               CollapseMixedReshapeOps<CollapseShapeOp, ExpandShapeOp>,
               CollapseShapeOpMemRefCastFolder>(context);
 }
 OpFoldResult ExpandShapeOp::fold(ArrayRef<Attribute> operands) {
-  if (succeeded(foldMemRefCast(*this)))
-    return getResult();
   return foldReshapeOp<ExpandShapeOp, CollapseShapeOp>(*this, operands);
 }
 OpFoldResult CollapseShapeOp::fold(ArrayRef<Attribute> operands) {
   return foldReshapeOp<CollapseShapeOp, ExpandShapeOp>(*this, operands);
 }
 
 //===----------------------------------------------------------------------===//
 // ReshapeOp

mlir/test/Dialect/MemRef/canonicalize.mlir

%1 = memref.collapse_shape %0 [[0, 1], [2]]
       : memref<?x4x?xf32> into memref<?x?xf32>
   return %1 : memref<?x?xf32>
 }
 // CHECK-LABEL: @fold_memref_reshape_dynamic
 //   CHECK-NOT:   linalg.{{.*}}_shape
 
 // -----
 
+func @fold_memref_expand_cast(%arg0 : memref<?x?xf32>) -> memref<2x4x4xf32> {
+  %0 = memref.cast %arg0 : memref<?x?xf32> to memref<8x4xf32>
+  %1 = memref.expand_shape %0 [[0, 1], [2]]
+      : memref<8x4xf32> into memref<2x4x4xf32>
+  return %1 : memref<2x4x4xf32>
+}
+
+// CHECK-LABEL: @fold_memref_expand_cast
+// CHECK: memref.expand_shape
+
+// -----
+
 // CHECK-LABEL:   func @collapse_after_memref_cast_type_change(
 // CHECK-SAME:      %[[INPUT:.*]]: memref<?x512x1x1xf32>) -> memref<?x?xf32> {
 // CHECK:           %[[COLLAPSED:.*]] = memref.collapse_shape %[[INPUT]]
 // CHECK-SAME:         {{\[\[}}0], [1, 2, 3]] : memref<?x512x1x1xf32> into memref<?x512xf32>
 // CHECK:           %[[DYNAMIC:.*]] = memref.cast %[[COLLAPSED]] :
 // CHECK-SAME:         memref<?x512xf32> to memref<?x?xf32>
 // CHECK:           return %[[DYNAMIC]] : memref<?x?xf32>
 // CHECK:         }