[Mlir-commits] [mlir] [MLIR] Add pattern to bubble up tensor.extract_slice (PR #126898)

[Andrzej Warzyński]

================
@@ -210,6 +214,217 @@ struct BubbleUpExpandThroughParallelCollapse
   }
 };
 
+/// Converts `tensor.extract_slice(tensor.expand_shape)` to
+/// `tensor.expand_shape(tensor.extract_slice)`.
+/// For this transformation to be possible, the slice must be fully contiguous
+/// within each reassociation group of the expand_shape.
+/// A slice is defined as fully contiguous within a reassociation group if after
+/// flattening the reassociation group to a single 1D range, then the slice
+/// taken out of the group could be defined as a single contiguous subrange
+/// within that range.
+/// If the transformation is not possible, or if the slice is rank reducing, the
+/// function returns failure.
+///
+/// Example:
+/// ```
+/// %reshape = tensor.expand_shape %in [[0, 1], [2, 3], [4, 5, 6]]
+///     tensor<8x16x32xf32> to tensor<2x4x2x8x4x2x4xf32>
+/// %slice = tensor.extract_slice %reshape ...
+///     tensor<2x4x2x8x4x2x4xf32> to tensor<2x4x1x5x1x1x4xf32>
+///
+/// // The transformation is possible because each reassociation group has a
+/// // contiguous slice. (i.e., [2x4->2x4], [2x8->1x5], [4x2x4->1x1x4])
+/// // After the transformation:
+///
+/// %slice = tensor.extract_slice %in ...
+///     tensor<8x16x32xf32> to tensor<8x5x4xf32>
+/// %reshape = tensor.expand_shape %slice [[0, 1], [2, 3], [4, 5, 6]]
+///     tensor<8x5x4xf32> to tensor<2x4x1x5x1x1x4xf32>
+/// ```
+///
+/// Note - this pattern could be reworked to be a swap pattern between
+/// `tensor.expand_shape` and `tensor.extract_slice`, but is currently
+/// implemented only as a bubble up pattern for `tensor.extract_slice`.
----------------
banach-space wrote:

> Am I missing something here in my understanding of the distinction between the types of patterns?

Your taxonomy makes sense to me, but I haven't seen that documented anywhere 😅 Nor any API's/examples of that. But there's a lot of corners of MLIR that I haven't touched!

My concern is that there's a non-zero chance that we are discussing something that will not make sense to somebody else reading this comment and who didn't participate in this discussion. In fact, your tests in "transform-op-fuse.mlir" are already making references to "swapping": `@swap_expand_shape_with_extract_slice`. So, somebody might wonder, what's missing to make this a "swap" pattern?

I think that you can disambiguate this a bit by replacing "reworked" with "extended":
>  Note - this pattern could be extended ...

https://github.com/llvm/llvm-project/pull/126898