Fix expand folder to avoid folding memref cast

[cathyzhyi]

Summary

memref.expand_shape has verification logic to make sure
src type is compatible with the result type.
Cast folding might add/remove static information for the src operand
of memref.expand_shape which will change operation to be invalid.
The only case where it's legal to fold the memref cast is actually
when the cast doesn't change the type info. This case is handled by
memref.cast folder already.

[silvas]

The idea here of preventing invalid ops makes sense to me, but @herhut probably is more knowledgeable about this code for a detailed review.

[herhut]

This was recently fixed in https://reviews.llvm.org/D114391. If you rebase, you should pick up that change.

Would be great to land those tests, though!

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

 // CHECK:           return %[[DYNAMIC]] : memref<?x?xf32>
 // CHECK:         }
 func @collapse_after_memref_cast_type_change(%arg0 : memref<?x512x1x1xf32>) -> memref<?x?xf32> {
   %dynamic = memref.cast %arg0: memref<?x512x1x1xf32> to memref<?x?x?x?xf32>
   %collapsed = memref.collapse_shape %dynamic [[0], [1, 2, 3]] : memref<?x?x?x?xf32> into memref<?x?xf32>
   return %collapsed : memref<?x?xf32>
 }
 
+// -----
+
 // CHECK-LABEL:   func @collapse_after_memref_cast(
 // CHECK-SAME:      %[[INPUT:.*]]: memref<?x512x1x?xf32>) -> memref<?x?xf32> {
 // CHECK:           %[[COLLAPSED:.*]] = memref.collapse_shape %[[INPUT]]
 // CHECK_SAME:        {{\[\[}}0], [1, 2, 3]] : memref<?x512x1x?xf32> into memref<?x?xf32>
 // CHECK:           return %[[COLLAPSED]] : memref<?x?xf32>
 func @collapse_after_memref_cast(%arg0 : memref<?x512x1x?xf32>) -> memref<?x?xf32> {
   %dynamic = memref.cast %arg0: memref<?x512x1x?xf32> to memref<?x?x?x?xf32>
   %collapsed = memref.collapse_shape %dynamic [[0], [1, 2, 3]] : memref<?x?x?x?xf32> into memref<?x?xf32>
   return %collapsed : memref<?x?xf32>
 }
+
+// -----
+
+// CHECK-LABEL:   func @expand_after_memref_cast_type_change(
+// CHECK-SAME:      %[[INPUT:.*]]: memref<1x?xf32>) -> memref<?x512x1x?xf32> {
+// CHECK:           %[[CAST:.*]] = memref.cast %[[INPUT]] : memref<1x?xf32> to memref<?x?xf32>
+// CHECK:           %[[EXPANDED:.*]] = memref.expand_shape %[[CAST]] {{\[\[}}0], [1, 2, 3]] : memref<?x?xf32> into memref<?x512x1x?xf32>
+// CHECK:           return %[[EXPANDED]] : memref<?x512x1x?xf32>
+func @expand_after_memref_cast_type_change(%arg0 : memref<1x?xf32>) -> memref<?x512x1x?xf32> {
+  %dynamic = memref.cast %arg0: memref<1x?xf32> to memref<?x?xf32>
+  %expanded = memref.expand_shape %dynamic [[0], [1, 2, 3]] : memref<?x?xf32> into memref<?x512x1x?xf32>
+  return %expanded: memref<?x512x1x?xf32>
+}
+
+// -----
+
+// CHECK-LABEL:   func @expand_after_memref_cast(
+// CHECK-SAME:      %[[INPUT:.*]]: memref<?x?xf32>) -> memref<?x512x1x?xf32> {
+// CHECK:           %[[EXPANDED:.*]] = memref.expand_shape %[[INPUT]] {{\[\[}}0], [1, 2, 3]] : memref<?x?xf32> into memref<?x512x1x?xf32>
+// CHECK:           return %[[EXPANDED]] : memref<?x512x1x?xf32>
+func @expand_after_memref_cast(%arg0 : memref<?x?xf32>) -> memref<?x512x1x?xf32> {
+  %dynamic = memref.cast %arg0: memref<?x?xf32> to memref<?x?xf32>
+  %expanded = memref.expand_shape %dynamic [[0], [1, 2, 3]] : memref<?x?xf32> into memref<?x512x1x?xf32>
+  return %expanded: memref<?x512x1x?xf32>
+}