[Mlir-commits] [mlir] [MLIR] Bubble up tensor.extract_slice through tensor.collapse_shape (PR #131982)

[Andrzej Warzyński]

================
@@ -428,6 +429,190 @@ struct BubbleUpExpandShapeThroughExtractSlice
   }
 };
 
+/// Converts `tensor.collapse_shape(tensor.extract_slice)` to
+/// `tensor.extract_slice(tensor.collapse_shape)`.
+///
+/// For this transformation to be possible, the slice must be representable as a
+/// contiguous slice within each reassociation group of the src.
----------------
banach-space wrote:

Going back to my example that's actually is extracting a contiguous slice:
```mlir
%extract_1 = tensor.extract_slice %src[0, 0, 0][1, 1, 10] : tensor<2x3x10xf32> to tensor <1x1x10>
%extract_2 = tensor.extract_slice %src[0, 0, 0][1, 2, 5] : tensor<2x3x10xf32> to tensor <1x1x5>
// Pseudo-code
%result = tensor.concat %extract_1,  %extract_2
```

IIUC, the actual condition for this to work is that _after_ bubbling up, the extraction of _the contiguous slice_ (*) is representable as a single slice obtained via `tensor.extract_slice`. You can use my example to demonstrate what would not qualify.

(*) Note, we are starting with a contiguous slice and even when using two `tensor.extract_slice` Ops, the actual final slice is also contiguous. It's not the lack of "contiguity" that's a problem here (that's always present), but the limitations of `tensor.extract_slice`.

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