Fix FoldReshapeOpWithUnitExtent when parentSrcType is rank-1

[asaadaldien]

Summary

When folding reshapes A -> B -> C and |A| = 1, the association map A -> C is a simple identity.
The current code doesn't handle this case so this diff add it explicitly.

[mravishankar]

mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp

407

Does this really need to be here. I think the logic below should cover this case as well. Maybe something is off there.

[asaadaldien]

mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp

407

Why not ? This is a trivial special case for reshapes A -> B -> C with rank(A) = 1 and result always have this form. Its fair to hoist trivial case and catch them up.
FWIW we can add a correctness assertion rank(C) = size(reassociationMaps). then generalize the code below in a case by case.

[mravishankar]

mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp

407

I looked into this and it seems like swapping the if condition at line 447 and line 450 below seems to fix the problem.

I really see this is a bug in the logic below, and unless its something substantial there is no point special casing this.

I have the fix (some changes to the test is needed as well). If you want I can send out the patch or you can update this one.

[asaadaldien]

mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp

407

Thanks! feel free to send the fix you already have it =). I will retire this diff.

However I still think doing this is harmless

if (cond_1) return cheap_cond_1_solution();

// Compute generic solution

assert(generic_solution);

[mravishankar]

mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp

407

Ok, I will send out the patch. If generic_solution handles cond_1 then why bother. If it doesnt then its not "generic" enough. Special casing might end up hiding bugs.

[asaadaldien]

mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp

407

Why bother -> cheaper to compute =).
These kind of transformations without fuzzing infrastructure hard to reason about correctness / coverage.

[mravishankar]

mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp

407

Can be retired in favor of D88521

Why bother -> cheaper to compute =).

These kind of transformations without fuzzing infrastructure hard to reason about correctness / coverage.

Still dont see whats the reason for special casing. If you can generalize the logic (as can be done here) then why not do it. If by cheaper you mean compile time cheaper, this is really negligible here.

[asaadaldien]

Abandoned instead https://reviews.llvm.org/D88521

mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp

LogicalResult matchAndRewrite(TensorReshapeOp reshapeOp,
     RankedTensorType srcType = reshapeOp.getSrcType(),
                      dstType = reshapeOp.getResultType(),
                      parentSrcType = parentReshapeOp.getSrcType();
     if (!srcType.hasStaticShape() || !dstType.hasStaticShape() ||
         !parentSrcType.hasStaticShape() ||
         srcType.getRank() < dstType.getRank() ||
         parentSrcType.getRank() == dstType.getRank())
       return failure();
+
+    // Rank-1 parentSrcType there will always have just an identity
+    // reassociation map.
+    if (parentSrcType.getRank() == 1) {
+      rewriter.replaceOpWithNewOp<TensorReshapeOp>(
+          reshapeOp, dstType, parentReshapeOp.src(),
+          rewriter.getAffineMapArrayAttr({AffineMap::getMultiDimIdentityMap(
+              dstType.getRank(), rewriter.getContext())}));
+      return success();
+    }
+
     // Check if the result tensor_reshape after folding the reshapeOp and
     // parentReshapeOp are combined.
     // If the final tensor_reshape is folding, the parentReshapeOp is
     // introducing unit-dims, and the reshapeOp does an actual reshape.
-    // If the final tensor_reshape op is expanding, the reshapeOp is introducing
-    // unit-dims, and the parentReshapeOp does an actual reshape.
+    // If the final tensor_reshape op is expanding, the reshapeOp is
+    // introducing unit-dims, and the parentReshapeOp does an actual reshape.
     bool isFoldingPattern = parentSrcType.getRank() > dstType.getRank();
     auto reassociationMaps = isFoldingPattern
                                  ? reshapeOp.getReassociationMaps()
                                  : parentReshapeOp.getReassociationMaps();
     DenseSet<unsigned> conservedDimensions;
     for (auto &map : reassociationMaps) {
       if (map.getNumResults() == 1) {
         conservedDimensions.insert(

mlir/test/Dialect/Linalg/drop-unit-extent-dims.mlir

%1 = linalg.tensor_reshape %0
     [affine_map<(d0, d1, d2, d3, d4, d5, d6, d7, d8) -> (d0, d1, d2)>,
      affine_map<(d0, d1, d2, d3, d4, d5, d6, d7, d8) -> (d3, d4)>,
      affine_map<(d0, d1, d2, d3, d4, d5, d6, d7, d8) -> (d5)>,
      affine_map<(d0, d1, d2, d3, d4, d5, d6, d7, d8) -> (d6, d7)>,
      affine_map<(d0, d1, d2, d3, d4, d5, d6, d7, d8) -> (d8)>]
     : tensor<1x4x1x512x1x1x512x1x4xf32> into tensor<4x512x1x512x4xf32>
   return %1 : tensor<4x512x1x512x4xf32>
 }
+
+// -----
+
+//   CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1) -> (d0, d1)>
+//       CHECK: func @fold_reshape
+//       CHECK: linalg.tensor_reshape %{{.*}} [#[[MAP0]]
+//  CHECK-SAME:   tensor<2xf32> into tensor<2x1xf32>
+func @fold_reshape(%arg0: tensor<2xf32>) -> tensor<2x1xf32>
+{
+  %0 = linalg.tensor_reshape %arg0 [affine_map<(d0, d1, d2) -> (d0, d1, d2)>] : tensor<2xf32> into tensor<2x1x1xf32>
+  %1 = linalg.tensor_reshape %0
+  [affine_map<(d0, d1, d2) -> (d0)>,
+   affine_map<(d0, d1, d2) -> (d1, d2)>
+  ] : tensor<2x1x1xf32> into tensor<2x1xf32>
+  return %1 : tensor<2x1xf32>
+}