[mlir][utils] Fix invalid reshapes in ComposeCollapseOfExpandOp

[springerm]

Summary

Do not generate CollapseShapeOps/ExpandShapeOps that have the same source and result shape. Generate a cast instead. Such reshapes became invalid with D138498.

mlir/include/mlir/Dialect/Utils/ReshapeOpsUtils.h

 ///     : tensor<?x?x?xi64> into tensor<?x?xi64>
 ///
 /// because [0] and [1] from `expand_shape` reassociation cover completely
 /// `[0, 1]` from `collapse_shape`. If it is impossible to find such union of
 /// indices, then we fail.
 //
 /// When `rank(srcType) < rank(resultType)`, then we just swap `reassociation_1`
 /// `reassociation_2` and produce `expand_shape`.
-template <typename CollapseOpTy, typename ExpandOpTy>
+template <typename CollapseOpTy, typename ExpandOpTy, typename CastOpTy>
 struct ComposeCollapseOfExpandOp : public OpRewritePattern<CollapseOpTy> {
   using OpRewritePattern<CollapseOpTy>::OpRewritePattern;
   LogicalResult matchAndRewrite(CollapseOpTy collapseOp,
                                 PatternRewriter &rewriter) const override {
     auto expandOp = collapseOp.getSrc().template getDefiningOp<ExpandOpTy>();
     if (!expandOp)
       return failure();
 
     ShapedType srcType = expandOp.getSrcType();
     ShapedType resultType = collapseOp.getResultType();
 
     if (hasNonIdentityLayout(collapseOp.getSrc().getType()) ||
         hasNonIdentityLayout(expandOp.getSrc().getType()) ||
         hasNonIdentityLayout(expandOp.getResult().getType()))
       return failure();
 
     int64_t srcRank = srcType.getRank();
     int64_t resultRank = resultType.getRank();
     if (srcType == resultType)
       return failure();
 
     SmallVector<ReassociationIndices, 4> higherRankReassociation,
         lowerRankReassociation;
 
-    bool isResultCollapsed = srcRank > resultRank;
-    if (isResultCollapsed) {
+    if (srcRank > resultRank) {
       higherRankReassociation = expandOp.getReassociationIndices();
       lowerRankReassociation = collapseOp.getReassociationIndices();
     } else {
       higherRankReassociation = collapseOp.getReassociationIndices();
       lowerRankReassociation = expandOp.getReassociationIndices();
     }
 
     size_t higherRankIndicesID = 0;
     SmallVector<ReassociationIndices, 4> composedReassociation;
     for (const auto &lowerRankIndices : lowerRankReassociation) {
       ReassociationIndices composedIndices;
       while (higherRankIndicesID < higherRankReassociation.size()) {
         auto rightmostIndex =
             higherRankReassociation[higherRankIndicesID].back();
         if (rightmostIndex > lowerRankIndices.back())
           return failure();
         composedIndices.push_back(higherRankIndicesID++);
         if (rightmostIndex == lowerRankIndices.back())
           break;
       }
       composedReassociation.push_back(composedIndices);
     }
-    if (isResultCollapsed)
+    if (srcRank > resultRank) {
       rewriter.replaceOpWithNewOp<CollapseOpTy>(
           collapseOp, resultType, expandOp.getSrc(), composedReassociation);
-    else
+    } else if (srcRank < resultRank) {
       rewriter.replaceOpWithNewOp<ExpandOpTy>(
           collapseOp, resultType, expandOp.getSrc(), composedReassociation);
+    } else {
+      // Collapses/expansions that do not change the rank are not allowed. Use
+      // a cast instead.
+      assert(llvm::equal(srcType.getShape(), resultType.getShape()) &&
+             "expected same shape");
+      rewriter.replaceOpWithNewOp<CastOpTy>(collapseOp, resultType,
+                                            expandOp.getSrc());
+    }
     return success();
   }
 };
 
 template <typename ExpandOpTy, typename CollapseOpTy>
 struct ComposeExpandOfCollapseOp : public OpRewritePattern<ExpandOpTy> {
   using OpRewritePattern<ExpandOpTy>::OpRewritePattern;
   LogicalResult matchAndRewrite(ExpandOpTy expandOp,

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

LogicalResult matchAndRewrite(CollapseShapeOp op,
     }
     return success();
   }
 };
 
 void CollapseShapeOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                                   MLIRContext *context) {
   results.add<ComposeReassociativeReshapeOps<CollapseShapeOp>,
-              ComposeCollapseOfExpandOp<CollapseShapeOp, ExpandShapeOp>,
+              ComposeCollapseOfExpandOp<CollapseShapeOp, ExpandShapeOp, CastOp>,
               CollapseShapeOpMemRefCastFolder>(context);
 }
 
 OpFoldResult ExpandShapeOp::fold(ArrayRef<Attribute> operands) {
   return foldReshapeOp<ExpandShapeOp, CollapseShapeOp>(*this, operands);
 }
 
 OpFoldResult CollapseShapeOp::fold(ArrayRef<Attribute> operands) {

mlir/lib/Dialect/Tensor/IR/TensorOps.cpp

results.add<ComposeReassociativeReshapeOps<ExpandShapeOp>,
               FoldReshapeWithFromElements<ExpandShapeOp>, FoldDimOfExpandShape,
               FoldDimOfCollapseShape>(context);
 }
 
 void CollapseShapeOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                                   MLIRContext *context) {
   results
       .add<ComposeReassociativeReshapeOps<CollapseShapeOp>,
-           ComposeCollapseOfExpandOp<CollapseShapeOp, ExpandShapeOp>,
+           ComposeCollapseOfExpandOp<CollapseShapeOp, ExpandShapeOp, CastOp>,
            FoldReshapeWithConstant<CollapseShapeOp>,
            FoldReshapeWithFromElements<CollapseShapeOp>, FoldCollapseOfCastOp>(
           context);
 }
 
 OpFoldResult ExpandShapeOp::fold(ArrayRef<Attribute> operands) {
   return foldReshapeOp<ExpandShapeOp, CollapseShapeOp>(*this, operands);
 }

mlir/test/Dialect/MemRef/canonicalize.mlir

 // CHECK-NOT: memref.realloc
 // CHECK: return %[[SRC]]
 func.func @memref_realloc_dead(%src : memref<2xf32>, %v : f32) -> memref<2xf32>{
   %0 = memref.realloc %src : memref<2xf32> to memref<4xf32>
   %i2 = arith.constant 2 : index
   memref.store %v, %0[%i2] : memref<4xf32>
   return %src : memref<2xf32>
 }
+
+// -----
+
+// CHECK-LABEL: func @collapse_expand_fold_to_cast(
+//  CHECK-SAME:     %[[m:.*]]: memref<?xf32, strided<[1]>, 3>
+//       CHECK:   %[[casted:.*]] = memref.cast %[[m]] : memref<?xf32, strided<[1]>, 3> to memref<?xf32, 3
+//       CHECK:   return %[[casted]]
+func.func @collapse_expand_fold_to_cast(%m: memref<?xf32, strided<[1]>, 3>)
+    -> (memref<?xf32, 3>)
+{
+  %0 = memref.expand_shape %m [[0, 1]]
+      : memref<?xf32, strided<[1]>, 3> into memref<1x?xf32, 3>
+  %1 = memref.collapse_shape %0 [[0, 1]]
+      : memref<1x?xf32, 3> into memref<?xf32, 3>
+  return %1 : memref<?xf32, 3>
+}

mlir/test/Dialect/Tensor/canonicalize.mlir

 //       CHECK:   return %[[apply]]
 func.func @dim_of_collapse_shape(%t: tensor<?x?x?x7x?xf32>) -> index {
   %c1 = arith.constant 1 : index
   %0 = tensor.collapse_shape %t [[0], [1, 2, 3, 4]]
       : tensor<?x?x?x7x?xf32> into tensor<?x?xf32>
   %1 = tensor.dim %0, %c1 : tensor<?x?xf32>
   return %1 : index
 }
+
+// -----
+
+// CHECK-LABEL: func @collapse_expand_fold_to_cast(
+//  CHECK-SAME:     %[[t:.*]]: tensor<?xf32>
+//       CHECK:   return %[[t]]
+func.func @collapse_expand_fold_to_cast(%t: tensor<?xf32>) -> (tensor<?xf32>)
+{
+  %0 = tensor.expand_shape %t [[0, 1]] : tensor<?xf32> into tensor<1x?xf32>
+  %1 = tensor.collapse_shape %0 [[0, 1]] : tensor<1x?xf32> into tensor<?xf32>
+  return %1 : tensor<?xf32>
+}