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

[Andrzej Warzyński]

================
@@ -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:
----------------
banach-space wrote:

I struggle with this part a bit. Shouldn't there be only one reassociation group in which slicing happens? And, within that group, exactly one expanded dim should be sliced? If I am incorrect, is there an example?

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