[fixup] Additional edge-case

Signed-off-by: Artem Gindinson <[email protected]>

Author: AGindinson

mlir/lib/Dialect/Utils/ReshapeOpsUtils.cpp

@@ -139,8 +139,14 @@ mlir::getReassociationIndicesForCollapse(ArrayRef<int64_t> sourceShape,
         return std::nullopt;
       continue;
     }
-    if (wasLastDimDynamic && isDynamic)
-      return std::nullopt;
+    // If the last 2 dimensions in the target were dynamic, the tail in the
+    // source shape cannot contain a dynamic value. E.g. ?x?->? is valid,
+    // however ?x?x10x?->?x? would be indeterminate.
+    if (wasLastDimDynamic && numTargetDims > 1 &&
+        targetShape[numTargetDims - 2] == ShapedType::kDynamic) {
+      if (isDynamic)
+        return std::nullopt;
+    }
     // If the last target dimension is static, only source dimensions of 1 are
     // acceptable.
     if (!wasLastDimDynamic && !isOne)

mlir/unittests/Dialect/Utils/ReshapeOpsUtilsTest.cpp

@@ -92,6 +92,10 @@ TEST(ReassociationIndicesForCollapse, DynamicTest) {
   EXPECT_EQ(getReassociationIndicesForCollapse({1, ShapedType::kDynamic, 1},
                                                {ShapedType::kDynamic}),
             makeOptionalIndices({{0, 1, 2}}));
+  EXPECT_EQ(getReassociationIndicesForCollapse(
+                {ShapedType::kDynamic, ShapedType::kDynamic, 1},
+                {ShapedType::kDynamic, ShapedType::kDynamic}),
+            makeOptionalIndices({{0}, {1, 2}}));
   EXPECT_EQ(getReassociationIndicesForCollapse(
                 {1, ShapedType::kDynamic, ShapedType::kDynamic},
                 {ShapedType::kDynamic, ShapedType::kDynamic}),
@@ -122,4 +126,9 @@ TEST(ReassociationIndicesForCollapse, DynamicTestFailure) {
           {ShapedType::kDynamic, ShapedType::kDynamic, ShapedType::kDynamic},
           {ShapedType::kDynamic, ShapedType::kDynamic}),
       std::nullopt);
+  EXPECT_EQ(getReassociationIndicesForCollapse(
+                {ShapedType::kDynamic, ShapedType::kDynamic, 10, 1,
+                 ShapedType::kDynamic},
+                {ShapedType::kDynamic, ShapedType::kDynamic}),
+            std::nullopt);
 }