[mlir][tensor] Fold rank-reducing insert_slice with inverse collapse_shape

[springerm]

Summary

Depends On D139103

[mravishankar]

Same concern from https://reviews.llvm.org/D139103 . Posting comment from there

My thumb rule is if an op is not a simple op -> op change it doesnt pay for itself as a canonicalization. I am happy to discuss if indeed it is a canonicalization, but forcing everything to be rank-reducing when canonicalization is invoked is a very opinionated solution. It should be opt-in (ftr, I am in favor of using rank-reduce slices everywhere, but not in favor of forcing that view on everyone)

[springerm]

rebase

[mravishankar]

This is yet another instance of a patch being submitted when it has still changes requested marked. Whats the hurry, and AFAIK is not in keeping with developer practices. When patches land this way it drops off the dashboard and I dont see the resolution of the issue. That aside, specifically here, https://reviews.llvm.org/D139118 was tagged as something that conflicted with this. Does it or does it not. I am interested to know whats happening here (and hence I requested changes). Please do not land patches with request changes unless its been a while from the reviewer.

mlir/lib/Dialect/Tensor/Transforms/ReshapePatterns.cpp

LogicalResult matchAndRewrite(ExpandShapeOp expandShapeOp,
     SmallVector<OpFoldResult> mixedSizes = extractSliceOp.getMixedSizes();
     SmallVector<OpFoldResult> mixedStrides = extractSliceOp.getMixedStrides();
     rewriter.replaceOpWithNewOp<tensor::ExtractSliceOp>(
         expandShapeOp, extractSliceOp.getSource(), mixedOffsets, mixedSizes,
         mixedStrides);
     return success();
   }
 };
+
+/// Fold insert_slice(collapse_shape) ops that cancel itself out.
+struct FoldInsertOfRankReducingInsert : public OpRewritePattern<InsertSliceOp> {
+  using OpRewritePattern<InsertSliceOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(InsertSliceOp insertSliceOp,
+                                PatternRewriter &rewriter) const override {
+    auto collapseShapeOp =
+        insertSliceOp.getSource().getDefiningOp<CollapseShapeOp>();
+    if (!collapseShapeOp)
+      return failure();
+    RankedTensorType srcType = collapseShapeOp.getSrcType();
+
+    // Only cases where the CollapseShapeOp can be folded away entirely are
+    // supported. Moreover, only simple cases where the resulting InsertSliceOp
+    // has no rank-reduction anymore are supported at the moment.
+    RankedTensorType nonReducingInsertType =
+        RankedTensorType::get(insertSliceOp.getStaticSizes(),
+                              insertSliceOp.getType().getElementType());
+    if (nonReducingInsertType != srcType)
+      return failure();
+
+    SmallVector<OpFoldResult> mixedOffsets = insertSliceOp.getMixedOffsets();
+    SmallVector<OpFoldResult> mixedSizes = insertSliceOp.getMixedSizes();
+    SmallVector<OpFoldResult> mixedStrides = insertSliceOp.getMixedStrides();
+    rewriter.replaceOpWithNewOp<tensor::InsertSliceOp>(
+        insertSliceOp, collapseShapeOp.getSrc(), insertSliceOp.getDest(),
+        mixedOffsets, mixedSizes, mixedStrides);
+    return success();
+  }
+};
 } // namespace
 
 void mlir::tensor::populateReassociativeReshapeFoldingPatterns(
     RewritePatternSet &patterns) {
-  patterns.add<FoldExpandOfRankReducingExtract>(patterns.getContext());
+  patterns.add<FoldExpandOfRankReducingExtract, FoldInsertOfRankReducingInsert>(
+      patterns.getContext());
 }

mlir/test/Dialect/Tensor/fold-reassociative-reshapes.mlir

func.func @expand_shape_of_rank_reducing_extract(
   %0 = tensor.extract_slice %t[0, 0, 0, 0][%idx, 1, 1, 5][1, 1, 1, 1]
       : tensor<?x?x?x?xf32> to tensor<?x1x5xf32>
   %1 = tensor.expand_shape %0 [[0], [1, 2], [3]]
       : tensor<?x1x5xf32> into tensor<?x1x1x5xf32>
   %2 = tensor.expand_shape %0 [[0, 1], [2], [3]]
       : tensor<?x1x5xf32> into tensor<?x1x1x5xf32>
   return %1, %2 : tensor<?x1x1x5xf32>, tensor<?x1x1x5xf32>
 }
+
+// -----
+
+// CHECK-LABEL: func @rank_reducing_insert_of_collapse_shape(
+//  CHECK-SAME:     %[[t:.*]]: tensor<?x1x1x5xf32>
+//       CHECK:   %[[insert:.*]] = tensor.insert_slice %[[t]] into %{{.*}}[0, 0, 0, 0] [%{{.*}}, 1, 1, 5] [1, 1, 1, 1] : tensor<?x1x1x5xf32> into tensor<?x?x?x?xf32>
+//       CHECK:   return %[[insert]]
+func.func @rank_reducing_insert_of_collapse_shape(
+    %t: tensor<?x1x1x5xf32>, %d: tensor<?x?x?x?xf32>, %sz: index)
+  -> tensor<?x?x?x?xf32> {
+  %0 = tensor.collapse_shape %t [[0, 1], [2], [3]]
+      : tensor<?x1x1x5xf32> into tensor<?x1x5xf32>
+  %1 = tensor.insert_slice %0 into %d[0, 0, 0, 0][%sz, 1, 1, 5][1, 1, 1, 1]
+      : tensor<?x1x5xf32> into tensor<?x?x?x?xf32>
+  return %1 : tensor<?x?x?x?xf32>
+}