[mlir][tensor] Do not fold rank-reduced extract_slice into unpack op.

[chelini]

Thanks!

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

struct FoldUnpackWithExtractSliceOp : public OpRewritePattern<ExtractSliceOp> {
   using OpRewritePattern<ExtractSliceOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(ExtractSliceOp sliceOp,
                                 PatternRewriter &rewriter) const override {
     auto unpackOp = sliceOp.getSource().getDefiningOp<UnPackOp>();
     if (!unpackOp)
       return failure();
 
+    if (sliceOp.getResultType().getRank() != unpackOp.getDestType().getRank()) {
+      return rewriter.notifyMatchFailure(
+          sliceOp, "rank-reduced folding is not supported");
+    }
+
     // Check all offsets are zeros, and all strides are ones.
     if (!areAllConstantIntValue(sliceOp.getMixedOffsets(), 0) ||
         !areAllConstantIntValue(sliceOp.getMixedStrides(), 1)) {
       return rewriter.notifyMatchFailure(
           sliceOp, "expects offsets to be 0s and strides to be 1s");
     }
 
     // Create a new empty output tensor.

mlir/test/Dialect/Tensor/fold-into-pack-and-unpack.mlir

func.func @nofold_unpack_slice_non_unit_stride(%arg0 : tensor<?x?x8x4xf32>, %arg1 : tensor<?x?xf32>,
   return %1 : tensor<?x?xf32>
 }
 // CHECK-LABEL: func @nofold_unpack_slice_non_unit_stride(
 //       CHECK:   %[[UNPACK:.+]] = tensor.unpack
 //       CHECK:   tensor.extract_slice %[[UNPACK]]
 
 // -----
 
+func.func @nofold_unpack_slice_rank_reduced(%arg0 : tensor<?x?x8x4xf32>, %arg1 : tensor<?x?xf32>,
+    %arg2 : index, %arg3 : index) -> tensor<f32> {
+  %0 = tensor.unpack %arg0 inner_dims_pos = [0, 1] inner_tiles = [8, 4] into %arg1
+      : tensor<?x?x8x4xf32> -> tensor<?x?xf32>
+  %1 = tensor.extract_slice %0[0, 0] [1, 1] [1, 1] : tensor<?x?xf32> to tensor<f32>
+  return %1 : tensor<f32>
+}
+// CHECK-LABEL: func @nofold_unpack_slice_rank_reduced(
+//       CHECK:   %[[UNPACK:.+]] = tensor.unpack
+//       CHECK:   tensor.extract_slice %[[UNPACK]]
+
+// -----
+
 func.func @pad_pack(%src: tensor<16641x16xf32>) -> tensor<2082x1x8x32xf32> {
   %c0 = arith.constant 0 : index
   %cst = arith.constant 0.000000e+00 : f32
   %padded = tensor.pad %src low[0, 0] high[15, 0] {
   ^bb0(%arg0: index, %arg1: index):
     tensor.yield %cst : f32
   } : tensor<16641x16xf32> to tensor<16656x16xf32>
   %empty = tensor.empty() : tensor<2082x1x8x32xf32>