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

[ofri frishman]

================
@@ -428,6 +429,239 @@ struct BubbleUpExpandShapeThroughExtractSlice
   }
 };
 
+/// Converts `tensor.extract_slice(tensor.collapse_shape)` to
+///          `tensor.collapse_shape(tensor.extract_slice)`.
+///
+/// For this transformation to be possible - after bubbling up, the extraction
+/// of the contiguous slice must be representable as a single slice obtained via
+/// tensor.extract_slice within each reassociation group of the src.
+///
+/// In case the size and offset extracted are static then this is possible if
+/// the following conditions are met within each reassociation group:
----------------
ofri-frishman wrote:

I'll start out with the second part of your question regarding each individual reassociation group.
We need to make sure that after we bubble up the `tensor.extract_slice` that the data we extract from the expanded tensor for this group is contiguous, since it was contiguous before the bubble up.
Now lets look at the post bubble up sizes [ 1, 1, ..., 1, Sk, Ak + 1, Ak + 2, ...,An ] and see why this makes sense.
If we start with the trailing dims, we want them all the be full, which means that the data is obviously contiguous. Then we have this one dim where we extract a slice that is smaller than the full size Sk, but things are still contiguous since we now have a contiguous slice of size Sk*(Prod(Ai) for i=k+1 to n). Note this this slice might not be with offset 0, but that is still fine since the original slice we extracted might not have been with offset 0. 
Then we have all those leading dime of size 1. Since they are size 1 this means that we could say that we are technically  slicing on this dim but in practice the size 1 means that we aren't breaking contiguity and just maybe changing the offset of the contiguous block that we extract.
to summarize - we could technically say that we are slicing on k different dims, but in practice the result is a single contiguous block of data as required.

Now regarding the first part of the question about different reassociation groups.
The data extracted from the collapsed tensor might have been sliced on multiple different collapsed dims which could make the entire original slice to not be contiguous, so it should be fine if slicing occurs on multiple different reassociation groups after the bubble up.
I'll give an example for this from the tests:
```
BEFORE:
%collapse = tensor.collapse_shape %src [[0, 1], [2, 3]] : tensor<6x5x3x10xf32> into tensor<30x30xf32>
%extract = tensor.extract_slice %collapse[5, 10][15, 10][1, 1] : tensor<30x30xf32> to tensor<15x10xf32>
AFTER:
%extracted_slice = tensor.extract_slice %arg0[1, 0, 1, 0] [3, 5, 1, 10] [1, 1, 1, 1] : tensor<6x5x3x10xf32> to tensor<3x5x1x10xf32>
%collapsed = tensor.collapse_shape %extracted_slice [[0, 1], [2, 3]] : tensor<3x5x1x10xf32> into tensor<15x10xf32>
```
In the example there are 2 different reassociation groups, where a slice is extracted from each one of them, and the result of the bubble up is still legal and the result of the `tensor.collapse_shape` post bubble up represents the same data as the result of the `tensor.extract_slice` pre bubble up.
We can also directly calculate the indices within the src tensor of the elements from the result and see that they are equal in both cases: [160:169], [190:199], [220:229]...

  

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