[fixup] Clarify some early-return cases

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

Author: AGindinson

mlir/lib/Dialect/Utils/ReshapeOpsUtils.cpp

@@ -293,6 +293,10 @@ mlir::getReassociationIndicesForCollapse(ArrayRef<int64_t> sourceShape,
                                          ArrayRef<int64_t> targetShape) {
   unsigned numSourceDims = sourceShape.size(),
            numTargetDims = targetShape.size();
+  // We're supposed to search for a collapsing reassociation. If the sizes
+  // match, there's no actual collapsing taking place - it's either a no-op or a
+  // `tensor.reshape`-style reassociation (that would be beyond the scope of
+  // this utility).
   if (numSourceDims <= numTargetDims)
     return std::nullopt;
   // Early handling for scalar target types.

mlir/unittests/Dialect/Utils/ReshapeOpsUtilsTest.cpp

@@ -54,11 +54,14 @@ TEST(ReassociationIndicesForCollapse, StaticTest) {
 }
 
 TEST(ReassociationIndicesForCollapse, StaticTestFailure) {
-  EXPECT_EQ(getReassociationIndicesForCollapse({10, 20}, {10}), std::nullopt);
+  // No-op reassociation
   EXPECT_EQ(getReassociationIndicesForCollapse({10, 20}, {10, 20}),
             std::nullopt);
+  // Invalid static reassociations
+  EXPECT_EQ(getReassociationIndicesForCollapse({10, 20}, {10}), std::nullopt);
   EXPECT_EQ(getReassociationIndicesForCollapse({10, 20, 30}, {200, 300}),
             std::nullopt);
+  // Non-collapsing (expanding) reassociation
   EXPECT_EQ(getReassociationIndicesForCollapse({10, 20, 30}, {1, 10, 20, 30}),
             std::nullopt);
 }