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

[Artem Gindinson]

AGindinson wrote:

@IanWood1, while making adjustments, I've realized:
1. For a variety of cases, the non-greedy traversal only works in one direction. Take `2x2x?x4x5x? into 2x?x20x?`: it's valid, but the reassociation would fail for R-to-L traversal without backtracking because of `?x2 into ?x2x2` at the end. `2x2x?x4x5x4x? into 2x?x20x?` would be a corresponding ambiguous input that we wouldn't be able to validate.
\
_**1b.**_ I've meddled with a same direction, greedy-ish algorithm as an alternative way to validate: only for subsequent dynamic shapes in the target, it would maximize the range towards the beginning. This allows to exclude subexpressions like `?x2x? into ?x?`, `?x?x? into ?x?` for the same direction of traversal, but doesn't help with your examples above. Actual greedy matching requires full-on backtracking...

2. ...which would essentially make this a **regex match problem**, because dynamic dimensions act as a `*` would. Best illustrated by palindrome cases like `2x1x2x...x?x...x2x1x2 into 2x?x2`, with the fine bonus of having to confirm the match as unambiguous (but allowing ambiguity for unit dimensions!) - per our discussions above.

For a full solution, I'm taking another quick look at an approach with slicing the shapes into mixed & static subshapes. IMO, full-on backtracking is an overkill given the context. If not, let's just stick with 1 and accept the missed cases - we'll still get a good increase in the number of supported cases. I'll post what I have EO tomorrow.

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