[Mlir-commits] [mlir] 8d474f1 - [mlir] Handle an edge case when folding reshapes with multiple trailing 1 dimensions

[Benjamin Kramer]

Author: Benjamin Kramer
Date: 2021-11-29T18:31:43+01:00
New Revision: 8d474f1d157530577f06ce3ef9187e1aaf31a59e

URL: https://github.com/llvm/llvm-project/commit/8d474f1d157530577f06ce3ef9187e1aaf31a59e
DIFF: https://github.com/llvm/llvm-project/commit/8d474f1d157530577f06ce3ef9187e1aaf31a59e.diff

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

We would exit early and miss this case.

Differential Revision: https://reviews.llvm.org/D114711

Added: 
    

Modified: 
    mlir/lib/Dialect/Utils/ReshapeOpsUtils.cpp
    mlir/test/Dialect/Linalg/canonicalize.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Utils/ReshapeOpsUtils.cpp b/mlir/lib/Dialect/Utils/ReshapeOpsUtils.cpp
index 919415a3f13ed..f1acfb3ab5042 100644
--- a/mlir/lib/Dialect/Utils/ReshapeOpsUtils.cpp
+++ b/mlir/lib/Dialect/Utils/ReshapeOpsUtils.cpp
@@ -73,6 +73,8 @@ mlir::getReassociationIndicesForReshape(ShapedType sourceType,
     // 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.

diff  --git a/mlir/test/Dialect/Linalg/canonicalize.mlir b/mlir/test/Dialect/Linalg/canonicalize.mlir
index e938f8c232e44..a64bd64749135 100644
--- a/mlir/test/Dialect/Linalg/canonicalize.mlir
+++ b/mlir/test/Dialect/Linalg/canonicalize.mlir
@@ -279,6 +279,20 @@ func @fold_reshape_trailing_unit_dims_dynamic(%arg0: tensor<1x1x?x1x1x1xf32>) ->
 
 // -----
 
+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]]