[mlir] Does not bubble up extract slice when it is rank-reducing.

[okkwon]

Summary

The bubble up logic was written by assuming the slice operation is
always a normal slice that outputs a tensor with the same rank.

[mravishankar]

Small suggestion, otherwise looks good!

mlir/lib/Dialect/Linalg/Transforms/BubbleUpExtractSlice.cpp

79

You can just check if the sliceop is rank reducing by using sliceOp.getType().getRank() != sliceOp.getSourceType().getRank().

[okkwon]

address the comment

[okkwon]

Thanks Mahesh. Much cleaner!

mlir/lib/Dialect/Linalg/Transforms/BubbleUpExtractSlice.cpp

LogicalResult matchAndRewrite(tensor::ExtractSliceOp sliceOp,
     if (!linalgOp.hasTensorSemantics()) {
       return rewriter.notifyMatchFailure(sliceOp,
                                          "expected tensor of linalg op");
     }
 
     if (!sliceOp.hasUnitStride())
       return rewriter.notifyMatchFailure(sliceOp, "expected unit stride");
 
+    if (sliceOp.getType().getRank() != sliceOp.getSourceType().getRank()) {
+      return rewriter.notifyMatchFailure(sliceOp, "expected no rank reduction");
+    }
+
     OpOperand *outOperand = linalgOp.getOutputOperand(0);
     AffineMap indexingMap = linalgOp.getTiedIndexingMap(outOperand);
     if (!indexingMap.isProjectedPermutation()) {
       return rewriter.notifyMatchFailure(
           sliceOp, "expected a projected permutation for output");
     }
 
     auto linalgLoc = linalgOp.getLoc();

mlir/test/Dialect/Linalg/bubble-up-extract-slice-op.mlir

 // CHECK: func @conv_slice
 // CHECK: %[[INIT:.+]] = linalg.init_tensor [1, 112, 112, 32] : tensor<1x112x112x32xf32>
 // CHECK: %[[SLICE0:.+]] = tensor.extract_slice %arg0[0, 128, 128, 0] [1, 65, 65, 3] [1, 1, 1, 1] : tensor<1x225x225x3xf32> to tensor<1x65x65x3xf32>
 // CHECK: %[[SLICE1:.+]] = tensor.extract_slice %arg1[0, 0, 0, 16] [3, 3, 3, 16] [1, 1, 1, 1] : tensor<3x3x3x32xf32> to tensor<3x3x3x16xf32>
 // CHECK: %[[SLICE2:.+]] = tensor.extract_slice %[[INIT]][0, 64, 64, 16] [1, 32, 32, 16] [1, 1, 1, 1] : tensor<1x112x112x32xf32> to tensor<1x32x32x16xf32>
 // CHECK: %[[FILL:.+]] = linalg.fill ins(%[[CST:.+]] : f32) outs(%[[SLICE2]] : tensor<1x32x32x16xf32>) -> tensor<1x32x32x16xf32>
 // CHECK: %[[CONV:.+]] = linalg.conv_2d_nhwc_hwcf {dilations = dense<1> : tensor<2xi64>, strides = dense<2> : tensor<2xi64>} ins(%[[SLICE0]], %[[SLICE1]] : tensor<1x65x65x3xf32>, tensor<3x3x3x16xf32>) outs(%[[FILL]] : tensor<1x32x32x16xf32>) -> tensor<1x32x32x16xf32>
 // CHECK: return %[[CONV]] : tensor<1x32x32x16xf32>
+
+//-----
+
+// The slice is not supposed to be bubbled up when it is rank-reducing.
+func @rank_reducing_slice(%width : index) -> tensor<1x1x1x?xf32> {
+  %cst = arith.constant 1.000000e+00 : f32
+  %init = linalg.init_tensor [1, %width] : tensor<1x?xf32>
+  %fill = linalg.fill ins(%cst : f32) outs(%init : tensor<1x?xf32>) -> tensor<1x?xf32>
+  %slice = tensor.extract_slice %fill[0, 0] [1, %width] [1, 1] : tensor<1x?xf32> to tensor<?xf32>
+  %expand = tensor.expand_shape %slice [[0, 1, 2, 3]] : tensor<?xf32> into tensor<1x1x1x?xf32>
+  return %expand : tensor<1x1x1x?xf32>
+}
+
+// CHECK: func @rank_reducing_slice
+// CHECK: %[[INIT:.+]] = linalg.init_tensor
+// CHECK: %[[FILL:.+]] = linalg.fill ins
+// CHECK: %[[SLICE:.+]] = tensor.extract_slice %[[FILL]]
+// CHECK: %[[EXPAND:.+]] = tensor.expand_shape %[[SLICE]]
+// CHECK: return %[[EXPAND]]