[fixup] Add scalar target tests & fix em

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

Author: AGindinson

mlir/lib/Dialect/Utils/ReshapeOpsUtils.cpp

@@ -292,7 +292,8 @@ mlir::getReassociationIndicesForCollapse(ArrayRef<int64_t> sourceShape,
     return std::nullopt;
   // Early handling for scalar target types.
   if (numTargetDims == 0) {
-    ReassociationIndices allSourceIndices(numSourceDims);
+    ReassociationIndices allSourceIndices;
+    allSourceIndices.reserve(numSourceDims);
     for (unsigned sourceDimIdx = 0; sourceDimIdx < numSourceDims;
          ++sourceDimIdx) {
       int64_t sourceSize = sourceShape[sourceDimIdx];

mlir/unittests/Dialect/Utils/ReshapeOpsUtilsTest.cpp

@@ -7,6 +7,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Utils/ReshapeOpsUtils.h"
+#include "mlir/IR/BuiltinTypeInterfaces.h"
 #include "llvm/ADT/STLExtras.h"
 #include "gtest/gtest.h"
 #include <optional>
@@ -20,6 +21,29 @@ makeOptionalIndices(std::initializer_list<ReassociationIndices> list) {
   return std::optional<SmallVector<ReassociationIndices>>(list);
 }
 
+TEST(ReassociationIndicesForCollapse, ScalarTest) {
+  EXPECT_EQ(getReassociationIndicesForCollapse({1}, {}),
+            makeOptionalIndices({{0}}));
+  EXPECT_EQ(getReassociationIndicesForCollapse({1, 1}, {}),
+            makeOptionalIndices({{0, 1}}));
+  EXPECT_EQ(getReassociationIndicesForCollapse({ShapedType::kDynamic}, {}),
+            makeOptionalIndices({{0}}));
+  EXPECT_EQ(getReassociationIndicesForCollapse({1, ShapedType::kDynamic,
+                                                ShapedType::kDynamic, 1,
+                                                ShapedType::kDynamic},
+                                               {}),
+            makeOptionalIndices({{0, 1, 2, 3, 4}}));
+}
+
+TEST(ReassociationIndicesForCollapse, ScalarTestFailure) {
+  EXPECT_EQ(getReassociationIndicesForCollapse({}, {}), std::nullopt);
+  EXPECT_EQ(getReassociationIndicesForCollapse({}, {1}), std::nullopt);
+  EXPECT_EQ(getReassociationIndicesForCollapse({2}, {}), std::nullopt);
+  EXPECT_EQ(
+      getReassociationIndicesForCollapse({1, 2, ShapedType::kDynamic, 1}, {}),
+      std::nullopt);
+}
+
 TEST(ReassociationIndicesForCollapse, StaticTest) {
   EXPECT_EQ(getReassociationIndicesForCollapse({10, 20}, {200}),
             makeOptionalIndices({{0, 1}}));