[mlir] Handle an edge case when folding reshapes with multiple trailing 1 dimensions

[bkramer]

Summary

We would exit early and miss this case.

mlir/lib/Dialect/Utils/ReshapeOpsUtils.cpp

while (sourceDim < sourceShape.size()) {
     if (prodOfCollapsedDims * sourceShape[sourceDim] != currTargetShape)
       return llvm::None;
 
     currIndices.push_back(sourceDim++);
     // If the reassociation is empty but the currIndices is not, this by
     // definition is folding unit-dimensions with the result being scalar type.
     // So only append the `currIndices` if reassociation map is not empty.
     if (targetDim == targetShape.size()) {
+      while (sourceDim < sourceShape.size())
+        currIndices.push_back(sourceDim++);
       if (!reassociationMap.empty() && !currIndices.empty())
         reassociationMap.back().append(currIndices.begin(), currIndices.end());
       // Break out of the loops. We should be done here.
       break;
     }
     reassociationMap.emplace_back(ReassociationIndices{});
     std::swap(reassociationMap.back(), currIndices);
     prodOfCollapsedDims = 1;

mlir/test/Dialect/Linalg/canonicalize.mlir

func @fold_reshape_trailing_unit_dims_dynamic(%arg0: tensor<1x1x?x1x1x1xf32>) -> tensor<?xf32>
   return %1 : tensor<?xf32>
 }
 //       CHECK: func @fold_reshape_trailing_unit_dims_dynamic
 //       CHECK: linalg.tensor_collapse_shape %{{.*}} {{\[}}[0, 1, 2, 3, 4, 5]]
 //  CHECK-SAME:   tensor<1x1x?x1x1x1xf32> into tensor<?xf32>
 
 // -----
 
+func @fold_reshape_trailing_unit_dims(%arg0: tensor<12x42x1x1xf32>) -> tensor<12x42xf32>
+{
+  %0 = linalg.tensor_expand_shape %arg0 [[0], [1], [2], [3, 4]]
+      : tensor<12x42x1x1xf32> into tensor<12x42x1x1x1xf32>
+  %1 = linalg.tensor_collapse_shape %0 [[0], [1, 2, 3, 4]]
+      : tensor<12x42x1x1x1xf32> into tensor<12x42xf32>
+  return %1 : tensor<12x42xf32>
+}
+//       CHECK: func @fold_reshape_trailing_unit_dims
+//       CHECK: linalg.tensor_collapse_shape %{{.*}} {{\[}}[0], [1, 2, 3]]
+//  CHECK-SAME:   tensor<12x42x1x1xf32> into tensor<12x42xf32>
+
+// -----
+
 func @no_fold_reshapes(%arg0 : tensor<?x?x?xf32>) -> tensor<?x?xf32>
 {
   %0 = linalg.tensor_expand_shape %arg0 [[0], [1], [2, 3]]
       : tensor<?x?x?xf32> into tensor<?x?x1x?xf32>
   %1 = linalg.tensor_collapse_shape %0 [[0], [1, 2, 3]]
       : tensor<?x?x1x?xf32> into tensor<?x?xf32>
   return %1 : tensor<?x?xf32>
 }