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

[Ian Wood]

https://github.com/IanWood1 requested changes to this pull request.

I think there is a bug here. For example:

```cpp
  EXPECT_EQ(getReassociationIndicesForCollapse(
                {ShapedType::kDynamic, 10, 10, 10, ShapedType::kDynamic},
                {ShapedType::kDynamic, 10, ShapedType::kDynamic}),
            std::nullopt);
```
This fails with the output `({ { 0 }, { 1 }, { 2, 3, 4 } })` (I think `std::nullopt` is correct result here)

Other examples:

```cpp
  EXPECT_EQ(getReassociationIndicesForCollapse(
                {ShapedType::kDynamic, 3, 4, 3, ShapedType::kDynamic},
                {ShapedType::kDynamic, 12, ShapedType::kDynamic}),
            std::nullopt);
  EXPECT_EQ(getReassociationIndicesForCollapse(
                {ShapedType::kDynamic, 8, 4, 2, 16, ShapedType::kDynamic},
                {ShapedType::kDynamic, 32, ShapedType::kDynamic}),
            std::nullopt);
  EXPECT_EQ(getReassociationIndicesForCollapse(
                {ShapedType::kDynamic, 2, 2, 2, ShapedType::kDynamic},
                {ShapedType::kDynamic, 2, 2, ShapedType::kDynamic}),
            std::nullopt);
```


I'm trying to think of a generalized rule for when it becomes ambiguous. One possible way might be to check against a reversed target/source shape:

```
getReassociationIndicesForCollapse(source, target) == 
	reverseReassociation(
		getReassociationIndicesForCollapse(reverse(source), reverse(target)))
```

The idea is that the current algorithm is "anti-greedy" and will try go minimize the size of the reassociations towards the beginning. So reversing the ordering will compare anti-greedy with greedy. The results will be the same if there is non-ambiguous reassociation.



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