[Mlir-commits] [mlir] [mlir][tensor] Loosen restrictions on folding dynamic reshapes (PR #137963)

[Artem Gindinson]

AGindinson wrote:

> Actual greedy matching requires full-on backtracking, which would essentially make this a regex match problem, <...>.

Too much drama, because:

> For a full solution, <...> an approach with slicing the shapes into mixed & static subshapes

This problem can be mathematically reduced to my current version. We just need to handle static head and/or tail separately from the dynamic/mixed subshape in between. A few examples for success & failure are below.

The logic from the current version would be used heavily, so a review iteration would be of great help already @IanWood1 @Mahesh. As you see fit, we can also push the current approach over the merge line, then I'll iterate further to a complete solution.

--

S1. `?x4x5x3x?x3x2 into ?x20x?x2`: For `?x4x5x3x? into ?x20x?` (mixed subshape), both the L-R and R-L lazy match yield `[0, 1, 2], [3, 4]`. For the tail slice of `3x2 into 2` we find the lazy match R-L, then assign the trailing `3` to collapse into the preceding dynamic dimension.
S2. `2x3x?x5x4x?x3x2 into 2x?x4x?x6`: the mixed subshapes of `?x5x4x? into ?x4x?` yield the same result L-R / R-L. We process the head of `2x3` into `2` L-R and collapse `3` dynamically. We match the tail of `3x2 into 6` R-to-L.
S3. `?x10x?x10 into ?x10`: the mixed subshapes are `?x10x? into ?`.

F1a. `2x1x2x?x4x5x4x? into 2x?x20x?`: For `?x4x5x4x? into ?x20x?` (mixed subshape), L-R and R-L lazy matches differ – no need to look at `2x1x2 into 2` (head).
F1b. `?x4x5x3x?x2x3x4x2 into ?x20x?x6`: as in S1, the mixed subshape is matched unambiguously. But R-L match fails for the tail of `2x3x4x2 into 6`.
F2a. `4x?x5 into 20`: mixed subshape in source, no mixed subshape in target
F2b. `2x?x1x1x?x5 into `?x?x?`: mixed subshape of the source (`?x1x1x?`) doesn't cover the "mixed subshape" of the fully dynamic target.
F3a. `?x10x?x10 into ?x?x10`: as opposed to S3, the mixed subshape reassociation (`?x10x? into ?x?`) will differ between L-R/R-L.
F3b. `?x?x? into ?x?`: The whole shape is the "mixed/dynamic subshape". L-R and R-L lazy matches differ.

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