[Mlir-commits] [mlir] [mlir] Fix ComposeExpandOfCollapseOp for dynamic case (PR #142663)

[Ian Wood]

================
@@ -387,11 +387,13 @@ struct ComposeExpandOfCollapseOp : public OpRewritePattern<ExpandOpTy> {
       auto resultSubShape =
           resultShape.slice(resultIndices.front(), resultIndices.size());
 
+      if (llvm::count_if(srcSubShape, ShapedType::isDynamic) >= 2)
+        return std::nullopt;
----------------
IanWood1 wrote:

Assuming the reassociation is

```
collapse: [[0, 1], [2], [3, 4, 5]]
expand: [[0], [1], [2, 3]]
```

The "subshape" reshape is `?x? -> ? -> ?` which is a trivial collapse but the check I added will skip it. Maybe both source and dst have to have >=2 dyn dims?

------


>Thinking out loud, but if there's more edge cases like this, maybe the correct way forward (on this affine map-ignoring path that we dislike) would be join the algorithms here and in getReassociationIndicesForCollapse(), so that we're always dealing with full information.

I think `findCollapsingReassociation` (this PR) and `getReassociationIndicesForCollapse` are answering different questions:

- `findCollapsingReassociation`: can this arbitrary reshape be simplified to a collapse?
- `getReassociationIndicesForCollapse`: given a source/target shape that is known to be a collapse, what is the reassociation?

The idea is that `findCollapsingReassociation` must determine that the reshape is representable as a collapse before using `getReassociationIndicesForCollapse`.

Using the example from the previous PR, `getReassociationIndicesForCollapse(?x8x128x? -> ?x128x?)` is `[[0, 1], [2]]` because you know the reshape is a collapse. But because `?x8x128x? -> ? -> ?x128x?` collapses into a single dimension you don't know that the output dynamic dimensions have any relation to the input dynamic dimensions (except for the product is 8x greater). At runtime it could be `10x8x128x10 -> 102400 -> 1x128x800` and there is no way to collapse `10x8x128x10` into `1x128x800`

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