This patch extends the `BubbleUpExtractSliceThroughCollapseShape`
pattern to handle cases where `tensor.extract_slice` has a dynamic
offset.
During tile and fuse transformations, it is common to encounter IR where
`tensor.extract_slice` operations appear after `tensor.collapse_shape`.
These patterns are used as cleanup transformations to canonicalize the
IR by bubbling up the slice operation before the reshape. This enables
further optimizations and simplifications downstream.
Previously, the pattern only handled:
1. Static offsets and sizes.
2. Dynamic sizes with a single non-unit expanded dimension.
This left a gap for additional common cases where we may have:
- Dynamic offsets with size == 1 (single element extraction).
- Size greater than 1 but the offset is computed dynamically.
Regarding the first case, It's always legal to perform such a
transformation.
Regarding the second case (i.e. dynamic offset with static
size > 1) we can guarantee contiguity by more restricted
conditions:
1. The innermost expanded dimension is divisible by the slice size
(ensures the slice fits within a single "row")
2. The offset is probably a multiple of the slice size (ensures the
slice starts at an aligned boundary)
For example, given:
- collapse_shape: tensor<6x10xf32> -> tensor<60xf32>
- extract_slice at offset with size
If 10 % size == 0 and offset % 5 == 0, the slice is guaranteed
contiguous. The transformation delinearizes the offset to get [row, col]
indices and extracts [1, size] from the expanded 6x10 shape.
The offset divisibility check uses affine analysis to statically verify
that affine expressions (e.g., `idx * 5`) are multiples of the slice
size.
e7715584c7