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

[ofri frishman]

================
@@ -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`.
----------------
ofri-frishman wrote:

Maybe I don't understand correctly the distinction between a "swap" pattern and a "bubble up" pattern.
I though that one of the main distinctions between them is that a swap pattern can change the order of ops in both directions and bubble up only supports a single direction.
Currently, the pattern I added only supports moving a `tensor.extract_slice` that has a `tensor.expand_shape` as a producer to be before the `tensor.expand_shape` and doesn't support the other way around. So, I saw it as a bubble up pattern.
Am I missing something here in my understanding of the distinction between the types of patterns?

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