[mlir] Rewrite canonicalization of collapse(expand) and expand(collapse) .

[pifon2a]

Summary

The original implementation does not take reassociation indices into account,
which leads to bugs.

[nicolasvasilache]

mlir/include/mlir/Dialect/Utils/ReshapeOpsUtils.h

235–236

You need to bail on non-identity layout for MemRefType atm.
You can use a simple templated hasNonIdentityLayout helper specialization to bail out early.

290

you need to bail on non-identity layout for MemRefType atm

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

1962–1963

I don't this this can be correct in the absence of explicit layout map handling.

2003–2004

I don't this this can be correct in the absence of explicit layout map handling.

[nicolasvasilache]

mlir/include/mlir/Dialect/Utils/ReshapeOpsUtils.h

257

The complexity of the double loop logic + case disjunction looks unnecessary to me.
You could just use two maps.

DenseMap<unsigned, unsigned> map1;
unsigned index = 0, group = 0;
for (auto &reassoc : op1.getReassociationIndices()) {
  for (unsigned i = 0, e = reassoc.size(); i < e; ++i)
    map1.insert(std::pair<unsigned, unsigned>(index++, group));
  ++group;
}
// Same for map2, this could become a separate reshape util by itself.

Then you can simply:

SmallVector<ReassociationIndices> composedReassociation(targetRank);
for (unsigned i = 0, e = sourceRank; i < e; ++i)
  composedReassociation[map2[map1[i]]].push_back(i);

This should work for all 4 cases (Collapse+Collapse, Collapse+Expand, Expand+Collapse, Expand+Expand).
Then you can just:

if (rankExpanding)
    rewriter.replaceOpWithNewOp<ExpandOpTy>(
        secondOp, resultType, firstOp.src(), composedReassociation);
else
    rewriter.replaceOpWithNewOp<CollapseOpTy>(
        secondOp, resultType, firstOp.src(), composedReassociation);

mlir/test/Dialect/MemRef/canonicalize.mlir

388–389

Please add a test case with memref + layout and make sure it does not canonicalize.
Once they land, we can reuse layout computation utils for expand/collapse and expand (haha) the behavior.

[pifon2a]

Address the comments.

[pifon2a]

Rewrite of the expand(collapse

mlir/include/mlir/Dialect/Utils/ReshapeOpsUtils.h

257

I am not sure how this approach would work for the cases, when the composition is impossible.

smth like 2x5xf32 -> [0, 1] collapse to 10xf32 -> [0, 1] expand to 5x2xf32.

Or ?x?xf32 -> [0], [1, 2] expand to ?x?x2xf32 -> [0, 1], [2] collapse to ?x2xf32.

[nicolasvasilache]

mlir/include/mlir/Dialect/Utils/ReshapeOpsUtils.h

257

You could factor out relevants parts of the verifier as to ensure the resulting Expand/Collapse is valid and bail on invalid rewrites.
We do similar type of probing in other places already (at least around cast ops areCastCompatible and some others I can't remember offhand).

[nicolasvasilache]

mlir/include/mlir/Dialect/Utils/ReshapeOpsUtils.h

257

Please disregard my suggestion, the intuition was wrong and too simplistic, my apologies for the noise.

Still, if there are ways to reuse parts of the verifier (and rewrite where necessary if affinemap is too annoying), please consider it rather than introducing new non-trivial code.

Thanks!

[pifon2a]

Update the pattern for expand(collapse).

[pifon2a]

mlir/include/mlir/Dialect/Utils/ReshapeOpsUtils.h

257

I reused parts of getReassociationIndicesForReshape.

[pifon2a]

Update the strided example.

[pifon2a]

Fix compiler warnings.

[bkramer]

Looks good, let's get this rolling to unblock stuff.

[pifon2a]

@nicolasvasilache I pushed the commit to fix the bugs. If you have further suggestions/comments, I ll be happy to address them.

[hanchung]

I think there are bugs in the commit. The verifier failed with the commit. To repro:

$ mlir-opt -canonicalize a.mlir

func @foo(%0: tensor<1x1xf32>, %1: tensor<1x1xf32>, %2: tensor<1x1xf32>) -> tensor<1x1xf32> {
  %cst = arith.constant 0.000000e+00 : f32
  %3 = linalg.init_tensor [8, 1] : tensor<8x1xf32>
  %4 = linalg.fill ins(%cst : f32) outs(%3 : tensor<8x1xf32>) -> tensor<8x1xf32>
  %5 = tensor.collapse_shape %0 [] : tensor<1x1xf32> into tensor<f32>
  %6 = tensor.insert_slice %5 into %4[0, 0] [1, 1] [1, 1] : tensor<f32> into tensor<8x1xf32>
  %7 = linalg.init_tensor [8, 1] : tensor<8x1xf32>
  %8 = linalg.fill ins(%cst : f32) outs(%7 : tensor<8x1xf32>) -> tensor<8x1xf32>
  %9 = tensor.collapse_shape %2 [] : tensor<1x1xf32> into tensor<f32>
  %10 = tensor.insert_slice %9 into %8[0, 0] [1, 1] [1, 1] : tensor<f32> into tensor<8x1xf32>
  %11 = tensor.collapse_shape %6 [[0, 1]] : tensor<8x1xf32> into tensor<8xf32>
  %12 = linalg.init_tensor [8] : tensor<8xf32>
  %13 = linalg.generic {indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>], iterator_types = ["parallel"]} ins(%11 : tensor<8xf32>) outs(%12 : tensor<8xf32>) {
  ^bb0(%arg3: f32, %arg4: f32):
    linalg.yield %arg3 : f32
  } -> tensor<8xf32>
  %14 = tensor.expand_shape %13 [[0, 1, 2, 3]] : tensor<8xf32> into tensor<1x1x8x1xf32>
  %15 = tensor.collapse_shape %1 [] : tensor<1x1xf32> into tensor<f32>
  %16 = linalg.init_tensor [] : tensor<f32>
  %17 = linalg.generic {indexing_maps = [affine_map<() -> ()>, affine_map<() -> ()>], iterator_types = []} ins(%15 : tensor<f32>) outs(%16 : tensor<f32>) {
  ^bb0(%arg3: f32, %arg4: f32):
    linalg.yield %arg3 : f32
  } -> tensor<f32>
  %18 = tensor.expand_shape %17 [] : tensor<f32> into tensor<1x1x1x1xf32>
  %19 = tensor.collapse_shape %10 [[0, 1]] : tensor<8x1xf32> into tensor<8xf32>
  %20 = linalg.init_tensor [8] : tensor<8xf32>
  %21 = linalg.generic {indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>], iterator_types = ["parallel"]} ins(%19 : tensor<8xf32>) outs(%20 : tensor<8xf32>) {
  ^bb0(%arg3: f32, %arg4: f32):
    linalg.yield %arg3 : f32
  } -> tensor<8xf32>
  %22 = tensor.expand_shape %21 [[0, 1, 2, 3]] : tensor<8xf32> into tensor<1x1x8x1xf32>
  %23 = linalg.mmt4d {comment = "f32*f32->f32, aarch64, matrix*vector"} ins(%14, %18 : tensor<1x1x8x1xf32>, tensor<1x1x1x1xf32>) outs(%22 : tensor<1x1x8x1xf32>) -> tensor<1x1x8x1xf32>
  %24 = tensor.collapse_shape %23 [[0, 1, 2, 3]] : tensor<1x1x8x1xf32> into tensor<8xf32>
  %25 = linalg.init_tensor [8] : tensor<8xf32>
  %26 = linalg.generic {indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>], iterator_types = ["parallel"]} ins(%24 : tensor<8xf32>) outs(%25 : tensor<8xf32>) {
  ^bb0(%arg3: f32, %arg4: f32):
    linalg.yield %arg3 : f32
  } -> tensor<8xf32>
  %27 = tensor.expand_shape %26 [[0, 1]] : tensor<8xf32> into tensor<8x1xf32>
  %28 = tensor.extract_slice %27[0, 0] [1, 1] [1, 1] : tensor<8x1xf32> to tensor<f32>
  %29 = tensor.expand_shape %28 [] : tensor<f32> into tensor<1x1xf32>
  return %29 : tensor<1x1xf32>
}

results in

a.mlir:24:9: error: 'tensor.expand_shape' op expected rank of the collapsed type(2) to be the number of reassociation maps(0)
  %18 = tensor.expand_shape %17 [] : tensor<f32> into tensor<1x1x1x1xf32>                                                                                                           ^
a.mlir:24:9: note: see current operation: %10 = "tensor.expand_shape"(%arg1) {reassociation = []} : (tensor<1x1xf32>) -> tensor<1x1x1x1xf32>

mlir/include/mlir/Dialect/Utils/ReshapeOpsUtils.h

SmallVector<ReassociationIndices, 2> convertReassociationMapsToIndices(
     OpBuilder &b, ArrayRef<ReassociationExprs> reassociationExprs);
 
 /// Return the reassociations maps to use to reshape given the source type and
 /// the target type when possible. Return llvm::None when this computation
 /// failed.
 Optional<SmallVector<ReassociationIndices>>
 getReassociationIndicesForReshape(ShapedType sourceType, ShapedType targetType);
 
+/// Returns the reassociation maps to collapse `sourceShape` to `targetShape` if
+/// possible.
+Optional<SmallVector<ReassociationIndices>>
+getReassociationIndicesForCollapse(ArrayRef<int64_t> sourceShape,
+                                   ArrayRef<int64_t> targetShape);
+
 /// Return true if the reassociation specification is valid, false otherwise.
 /// When false, the `invalidIndex` integer pointer is optionally filled with the
 /// index of the offending reassociation map.
 bool isReassociationValid(ArrayRef<AffineMap> reassociation,
                           int *invalidIndex = nullptr);
 
 template <typename ReshapeOpTy, typename InverseReshapeOpTy>
 static OpFoldResult foldReshapeOp(ReshapeOpTy reshapeOp,
static LogicalResult verifyReshapeLikeShapes(OpTy op, ShapedType collapsedType,
                                              ShapedType expandedType,
                                              bool isExpandingReshape) {
   return reshapeLikeShapesAreCompatible(
       [&](const Twine &msg) { return op->emitOpError(msg); },
       collapsedType.getShape(), expandedType.getShape(),
       op.getReassociationIndices(), isExpandingReshape);
 }
 
+/// Returns true iff the type is a MemRefType and has a non-identity layout.
+bool hasNonIdentityLayout(Type type);
+
 /// Pattern to collapse producer/consumer reshape ops that are both collapsing
 /// dimensions or are both expanding dimensions.
 template <typename ReshapeOpTy>
-struct CollapseReshapeOps : public OpRewritePattern<ReshapeOpTy> {
+struct ComposeReassociativeReshapeOps : public OpRewritePattern<ReshapeOpTy> {
   using OpRewritePattern<ReshapeOpTy>::OpRewritePattern;
   LogicalResult matchAndRewrite(ReshapeOpTy reshapeOp,
                                 PatternRewriter &rewriter) const override {
     auto srcReshapeOp = reshapeOp.src().template getDefiningOp<ReshapeOpTy>();
     if (!srcReshapeOp)
       return failure();
 
     ShapedType resultType = reshapeOp.getResultType();
+
+    if (hasNonIdentityLayout(srcReshapeOp.src().getType()) ||
+        hasNonIdentityLayout(reshapeOp.src().getType()) ||
+        hasNonIdentityLayout(reshapeOp.result().getType()))
+      return failure();
+
     Optional<SmallVector<ReassociationIndices>> reassociationIndices =
         composeReassociationIndices(srcReshapeOp.getReassociationIndices(),
                                     reshapeOp.getReassociationIndices(),
                                     rewriter.getContext());
     if (!reassociationIndices)
       return failure();
     rewriter.replaceOpWithNewOp<ReshapeOpTy>(
         reshapeOp, resultType, srcReshapeOp.src(), *reassociationIndices);
     return success();
   }
 };
 
-/// Pattern to collapse producer/consumer reshape ops that are both collapsing
-/// dimensions or are both expanding dimensions.
-template <typename ReshapeOpTy, typename InverseReshapeOpTy>
-struct CollapseMixedReshapeOps : public OpRewritePattern<ReshapeOpTy> {
-  using OpRewritePattern<ReshapeOpTy>::OpRewritePattern;
-  LogicalResult matchAndRewrite(ReshapeOpTy reshapeOp,
+/// Pattern to compose
+/// `collapse_shape(expand_shape(%src, reassociation_1), reassociation_2)`.
+/// In that case both `srcType` and `resultType` can be expressed as a function
+/// of `intermediateType`.
+/// In order to demonstrate the approach, let's assume that `rank(srcType) >
+/// `rank(resultType)`, i.e. the resulting operation should be `collapse_shape`.
+/// In that case, we can iterate over every set of indices in `reassociation_2`
+/// and try to find ids of sets of indices in `reassociation_1` that cover it
+/// completely.
+///
+/// Example:
+///
+///   %0 = tensor.expand_shape %arg [[0], [1], [2, 3]]
+///     : tensor<?x?x?xi64> into tensor<?x?x?x1xi64>
+///   %1 = tensor.collapse_shape %0 [[0, 1], [2, 3]]
+///     : tensor<?x?x?x1xi64> into tensor<?x?xi64>
+///
+/// can be canonicalized into
+///
+///   %0 = tensor.collapse_shape %arg [[0, 1], [2]]
+///     : tensor<?x?x?xi64> into tensor<?x?xi64>
+///
+/// because [0] and [1] from `expand_shape` reassociation cover completely
+/// `[0, 1]` from `collapse_shape`. If it is impossible to find such union of
+/// indices, then we fail.
+//
+/// When `rank(srcType) < rank(resultType)`, then we just swap `reassociation_1`
+/// `reassociation_2` and produce `expand_shape`.
+template <typename CollapseOpTy, typename ExpandOpTy>
+struct ComposeCollapseOfExpandOp : public OpRewritePattern<CollapseOpTy> {
+  using OpRewritePattern<CollapseOpTy>::OpRewritePattern;
+  LogicalResult matchAndRewrite(CollapseOpTy collapseOp,
                                 PatternRewriter &rewriter) const override {
-    auto srcReshapeOp =
-        reshapeOp.src().template getDefiningOp<InverseReshapeOpTy>();
-    if (!srcReshapeOp)
+    auto expandOp = collapseOp.src().template getDefiningOp<ExpandOpTy>();
+    if (!expandOp)
       return failure();
 
-    ShapedType srcReshapeSrcType = srcReshapeOp.getSrcType();
-    ShapedType intermediateType = reshapeOp.getSrcType();
-    ShapedType resultType = reshapeOp.getResultType();
+    ShapedType srcType = expandOp.getSrcType();
+    ShapedType resultType = collapseOp.getResultType();
 
-    // If the source reshape can be collapsed/expanded into the target reshape
-    // they can still be folded. This can only be reasoned about statically
-    // for cases where
-    // - either all shapes are static, or
-    // - The number of dynamic dimensions matches in the source of source and
-    //   result with all other dimensions being 1.
-    Optional<SmallVector<ReassociationIndices>> reassociationIndices =
-        getReassociationIndicesForReshape(srcReshapeSrcType, resultType);
-    if (!reassociationIndices)
+    if (hasNonIdentityLayout(collapseOp.src().getType()) ||
+        hasNonIdentityLayout(expandOp.src().getType()) ||
+        hasNonIdentityLayout(expandOp.result().getType()))
       return failure();
-    bool originalOpExpands =
-        intermediateType.getRank() > srcReshapeSrcType.getRank();
-    bool resultingOpExpands =
-        resultType.getRank() > srcReshapeSrcType.getRank();
-    if (!(resultingOpExpands ^ originalOpExpands))
-      rewriter.replaceOpWithNewOp<InverseReshapeOpTy>(
-          reshapeOp, resultType, srcReshapeOp.src(), *reassociationIndices);
+
+    int64_t srcRank = srcType.getRank();
+    int64_t resultRank = resultType.getRank();
+    if (srcType == resultType)
+      return failure();
+
+    SmallVector<ReassociationIndices, 4> higherRankReassociation,
+        lowerRankReassociation;
+
+    bool isResultCollapsed = srcRank > resultRank;
+    if (isResultCollapsed) {
+      higherRankReassociation = expandOp.getReassociationIndices();
+      lowerRankReassociation = collapseOp.getReassociationIndices();
+    } else {
+      higherRankReassociation = collapseOp.getReassociationIndices();
+      lowerRankReassociation = expandOp.getReassociationIndices();
+    }
+
+    size_t higherRankIndicesID = 0;
+    SmallVector<ReassociationIndices, 4> composedReassociation;
+    for (const auto &lowerRankIndices : lowerRankReassociation) {
+      ReassociationIndices composedIndices;
+      while (higherRankIndicesID < higherRankReassociation.size()) {
+        auto rightmostIndex =
+            higherRankReassociation[higherRankIndicesID].back();
+        if (rightmostIndex > lowerRankIndices.back())
+          return failure();
+        composedIndices.push_back(higherRankIndicesID++);
+        if (rightmostIndex == lowerRankIndices.back())
+          break;
+      }
+      composedReassociation.push_back(composedIndices);
+    }
+    if (isResultCollapsed)
+      rewriter.replaceOpWithNewOp<CollapseOpTy>(
+          collapseOp, resultType, expandOp.src(), composedReassociation);
     else
-      rewriter.replaceOpWithNewOp<ReshapeOpTy>(
-          reshapeOp, resultType, srcReshapeOp.src(), *reassociationIndices);
+      rewriter.replaceOpWithNewOp<ExpandOpTy>(
+          collapseOp, resultType, expandOp.src(), composedReassociation);
     return success();
   }
 };
 
+template <typename ExpandOpTy, typename CollapseOpTy>
+struct ComposeExpandOfCollapseOp : public OpRewritePattern<ExpandOpTy> {
+  using OpRewritePattern<ExpandOpTy>::OpRewritePattern;
+  LogicalResult matchAndRewrite(ExpandOpTy expandOp,
+                                PatternRewriter &rewriter) const override {
+    auto collapseOp = expandOp.src().template getDefiningOp<CollapseOpTy>();
+    if (!collapseOp)
+      return failure();
+
+    ShapedType srcType = collapseOp.getSrcType();
+    ShapedType resultType = expandOp.getResultType();
+
+    if (hasNonIdentityLayout(expandOp.src().getType()) ||
+        hasNonIdentityLayout(collapseOp.src().getType()) ||
+        hasNonIdentityLayout(collapseOp.result().getType()))
+      return failure();
+
+    int64_t srcRank = srcType.getRank();
+    int64_t resultRank = resultType.getRank();
+    if (srcType == resultType)
+      return failure();
+
+    auto srcReassociation = collapseOp.getReassociationIndices();
+    auto resultReassociation = expandOp.getReassociationIndices();
+    if (srcRank > resultRank) {
+      auto composedReassociation = findCollapsingReassociation(
+          srcReassociation, resultReassociation, srcType.getShape(),
+          resultType.getShape());
+      if (!composedReassociation.hasValue())
+        return failure();
+
+      rewriter.replaceOpWithNewOp<CollapseOpTy>(
+          expandOp, resultType, collapseOp.src(), *composedReassociation);
+      return success();
+    }
+    auto composedReassociation =
+        findCollapsingReassociation(resultReassociation, srcReassociation,
+                                    resultType.getShape(), srcType.getShape());
+    if (!composedReassociation.hasValue())
+      return failure();
+
+    rewriter.replaceOpWithNewOp<ExpandOpTy>(
+        expandOp, resultType, collapseOp.src(), *composedReassociation);
+    return success();
+  }
+
+private:
+  // Attempts to find a way to collapse `srcShape` to `resultShape` by
+  // collapsing subshapes defined by the reassociation indices.
+  Optional<SmallVector<ReassociationIndices>> findCollapsingReassociation(
+      ArrayRef<ReassociationIndices> srcReassociation,
+      ArrayRef<ReassociationIndices> resultReassociation,
+      ArrayRef<int64_t> srcShape, ArrayRef<int64_t> resultShape) const {
+    SmallVector<ReassociationIndices, 4> composedReassociation;
+
+    for (auto item : llvm::zip(srcReassociation, resultReassociation)) {
+      auto &srcIndices = std::get<0>(item);
+      auto &resultIndices = std::get<1>(item);
+      auto srcSubShape = srcShape.slice(srcIndices.front(), srcIndices.size());
+      auto resultSubShape =
+          resultShape.slice(resultIndices.front(), resultIndices.size());
+
+      if (srcSubShape.size() == resultSubShape.size()) {
+        if (srcSubShape == resultSubShape)
+          composedReassociation.push_back(srcIndices);
+        else
+          return llvm::None;
+      }
+
+      // Find reassociation to collapse `srcSubShape` into `resultSubShape`.
+      auto subShapeReassociation =
+          getReassociationIndicesForCollapse(srcSubShape, resultSubShape);
+      if (!subShapeReassociation.hasValue())
+        return llvm::None;
+
+      // Remap the subshape indices back to the original srcShape.
+      for (auto &subshape_indices : *subShapeReassociation) {
+        ReassociationIndices shape_indices;
+        for (int64_t index : subshape_indices)
+          shape_indices.push_back(srcIndices.front() + index);
+        composedReassociation.push_back(shape_indices);
+      }
+    }
+    return {std::move(composedReassociation)};
+  }
+};
+
 } // namespace mlir
 
 #endif // MLIR_DIALECT_UTILS_RESHAPEOPSUTILS_H

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

if (*expectedResultType != canonicalizedResultType)
     return emitOpError("expected expanded type to be ")
            << *expectedResultType << " but found " << canonicalizedResultType;
 
   return success();
 }
 
 void ExpandShapeOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                                 MLIRContext *context) {
-  results.add<CollapseReshapeOps<ExpandShapeOp>,
-              CollapseMixedReshapeOps<ExpandShapeOp, CollapseShapeOp>>(context);
+  results.add<ComposeReassociativeReshapeOps<ExpandShapeOp>,
+              ComposeExpandOfCollapseOp<ExpandShapeOp, CollapseShapeOp>>(
+      context);
 }
 
 /// Compute the layout map after collapsing a given source MemRef type with the
 /// specified reassociation indices.
 ///
 /// Note: All collapsed dims in a reassociation group must be contiguous. It is
 /// not possible to check this by inspecting a MemRefType in the general case.
 /// But it is assumed. If this is not the case, the behavior is undefined.
if (srcType.getLayout().isIdentity()) {
     if (failed(computedLayout))
       return emitOpError(
           "invalid source layout map or collapsing non-contiguous dims");
     auto computedType =
         MemRefType::get(resultType.getShape(), srcType.getElementType(),
                         *computedLayout, srcType.getMemorySpaceAsInt());
     expectedResultType = canonicalizeStridedLayout(computedType);
   }
 
   auto canonicalizedResultType = canonicalizeStridedLayout(resultType);
   if (expectedResultType != canonicalizedResultType)
     return emitOpError("expected collapsed type to be ")
            << expectedResultType << " but found " << canonicalizedResultType;
 
   return success();
 }
 
 struct CollapseShapeOpMemRefCastFolder
if (newResultType == op.getResultType()) {
       rewriter.replaceOpWithNewOp<CastOp>(op, op.getType(), newOp);
     }
     return success();
   }
 };
 
 void CollapseShapeOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                                   MLIRContext *context) {
-  results.add<CollapseReshapeOps<CollapseShapeOp>,
-              CollapseMixedReshapeOps<CollapseShapeOp, ExpandShapeOp>,
+  results.add<ComposeReassociativeReshapeOps<CollapseShapeOp>,
+              ComposeCollapseOfExpandOp<CollapseShapeOp, ExpandShapeOp>,
               CollapseShapeOpMemRefCastFolder>(context);
 }
 
 OpFoldResult ExpandShapeOp::fold(ArrayRef<Attribute> operands) {
   return foldReshapeOp<ExpandShapeOp, CollapseShapeOp>(*this, operands);
 }
 
 OpFoldResult CollapseShapeOp::fold(ArrayRef<Attribute> operands) {

mlir/lib/Dialect/Tensor/IR/TensorOps.cpp

LogicalResult matchAndRewrite(TensorReshapeOp reshapeOp,
     return success();
   }
 };
 
 } // namespace
 
 void ExpandShapeOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                                 MLIRContext *context) {
-  results.add<CollapseReshapeOps<ExpandShapeOp>,
-              CollapseMixedReshapeOps<ExpandShapeOp, CollapseShapeOp>,
+  results.add<ComposeReassociativeReshapeOps<ExpandShapeOp>,
+              ComposeExpandOfCollapseOp<ExpandShapeOp, CollapseShapeOp>,
               FoldReshapeWithConstant<ExpandShapeOp>,
               FoldReshapeWithFromElements<ExpandShapeOp>>(context);
 }
 
 void CollapseShapeOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                                   MLIRContext *context) {
-  results.add<CollapseReshapeOps<CollapseShapeOp>,
-              CollapseMixedReshapeOps<CollapseShapeOp, ExpandShapeOp>,
+  results.add<ComposeReassociativeReshapeOps<CollapseShapeOp>,
+              ComposeCollapseOfExpandOp<CollapseShapeOp, ExpandShapeOp>,
               FoldReshapeWithConstant<CollapseShapeOp>,
               FoldReshapeWithFromElements<CollapseShapeOp>>(context);
 }
 
 OpFoldResult ExpandShapeOp::fold(ArrayRef<Attribute> operands) {
   return foldReshapeOp<ExpandShapeOp, CollapseShapeOp>(*this, operands);
 }
 OpFoldResult CollapseShapeOp::fold(ArrayRef<Attribute> operands) {

mlir/lib/Dialect/Utils/ReshapeOpsUtils.cpp

 
 #include <numeric>
 
 using namespace mlir;
 
 Optional<SmallVector<ReassociationIndices>>
 mlir::getReassociationIndicesForReshape(ShapedType sourceType,
                                         ShapedType targetType) {
-  // Make the sourceType greater rank than the targetType. If they are same
-  // rank, then its an unsupported reshape op.
-  if (sourceType.getRank() == targetType.getRank())
-    return llvm::None;
+  if (sourceType.getRank() > targetType.getRank())
+    return getReassociationIndicesForCollapse(sourceType.getShape(),
+                                              targetType.getShape());
   if (sourceType.getRank() < targetType.getRank())
-    std::swap(sourceType, targetType);
+    return getReassociationIndicesForCollapse(targetType.getShape(),
+                                              sourceType.getShape());
+  return llvm::None;
+}
 
-  ArrayRef<int64_t> sourceShape = sourceType.getShape();
-  ArrayRef<int64_t> targetShape = targetType.getShape();
+Optional<SmallVector<ReassociationIndices>>
+mlir::getReassociationIndicesForCollapse(ArrayRef<int64_t> sourceShape,
+                                         ArrayRef<int64_t> targetShape) {
+  if (sourceShape.size() <= targetShape.size())
+    return llvm::None;
   unsigned sourceDim = 0;
   SmallVector<ReassociationIndices> reassociationMap;
-  reassociationMap.reserve(targetType.getRank());
+  reassociationMap.reserve(targetShape.size());
 
   ReassociationIndices currIndices;
   int64_t prodOfCollapsedDims = 1;
   while (sourceDim < sourceShape.size()) {
     unsigned targetDim = reassociationMap.size();
     // If we have mapped all the target dimensions stop and handle the remaining
     // tail of size-1 dimensions explictly.
-    if (targetDim == targetType.getRank())
+    if (targetDim == targetShape.size())
       break;
 
     int64_t currTargetShape = targetShape[targetDim];
     while (sourceShape[sourceDim] != ShapedType::kDynamicSize &&
            prodOfCollapsedDims * sourceShape[sourceDim] < currTargetShape &&
            sourceDim < sourceShape.size()) {
       prodOfCollapsedDims *= sourceShape[sourceDim];
       currIndices.push_back(sourceDim++);
mlir::getSymbolLessAffineMaps(ArrayRef<ReassociationExprs> reassociation) {
   SmallVector<AffineMap, 4> maps;
   maps.reserve(reassociation.size());
   for (const auto &exprs : reassociation) {
     assert(!exprs.empty());
     maps.push_back(AffineMap::get(maxDim + 1, 0, exprs, exprs[0].getContext()));
   }
   return maps;
 }
+
 bool mlir::isReassociationValid(ArrayRef<AffineMap> reassociation,
                                 int *invalidIndex) {
   if (reassociation.empty())
     return true;
   unsigned nDims = reassociation[0].getNumDims();
   unsigned nextExpectedDim = 0;
   for (const auto &it : llvm::enumerate(reassociation)) {
     auto m = it.value();
if (dynamicShape) {
                          " of collapsed type to be static value of " +
                          Twine(linearizedStaticShape));
       }
     }
     expandedDimStart += map.value().size();
   }
   return success();
 }
+
+bool mlir::hasNonIdentityLayout(Type type) {
+  if (auto memrefType = type.dyn_cast<MemRefType>())
+    return !memrefType.getLayout().isIdentity();
+  return false;
+}

mlir/test/Dialect/MemRef/canonicalize.mlir

 func @allocator(%arg0 : memref<memref<?xi32>>, %arg1 : index)  {
   %0 = memref.alloc(%arg1) : memref<?xi32>
   memref.store %0, %arg0[] : memref<memref<?xi32>>
   return
 }
 
 // -----
 
-func @collapsing_memref_reshapes_to_zero_dim(%arg0 : memref<1x1x1xf32>)
-                                             -> memref<f32> {
+func @compose_collapse_of_collapse_zero_dim(%arg0 : memref<1x1x1xf32>)
   %0 = memref.collapse_shape %arg0 [[0, 1, 2]]
       : memref<1x1x1xf32> into memref<1xf32>
   %1 = memref.collapse_shape %0 [] : memref<1xf32> into memref<f32>
   return %1 : memref<f32>
 }
-// CHECK-LABEL: collapsing_memref_reshapes_to_zero
+// CHECK-LABEL: func @compose_collapse_of_collapse_zero_dim
 //       CHECK:   memref.collapse_shape %{{.*}} []
 //  CHECK-SAME:     memref<1x1x1xf32> into memref<f32>
 
 // -----
 
-func @collapsing_memref_reshapes(%arg0 : memref<?x?x?x?x?xf32>)
+func @compose_collapse_of_collapse(%arg0 : memref<?x?x?x?x?xf32>)
     -> memref<?x?xf32> {
   %0 = memref.collapse_shape %arg0 [[0, 1], [2], [3, 4]]
       : memref<?x?x?x?x?xf32> into memref<?x?x?xf32>
   %1 = memref.collapse_shape %0 [[0, 1], [2]]
       : memref<?x?x?xf32> into memref<?x?xf32>
   return %1 : memref<?x?xf32>
 }
-// CHECK-LABEL: collapsing_memref_reshapes
+// CHECK-LABEL: func @compose_collapse_of_collapse
 //       CHECK:   memref.collapse_shape %{{.*}} {{\[}}[0, 1, 2], [3, 4]]
 //   CHECK-NOT:   memref.collapse_shape
 
 // -----
 
-func @expanding_memref_reshapes(%arg0 : memref<?x?xf32>)
+func @do_not_compose_collapse_of_expand_non_identity_layout(
+    %arg0: memref<?x?xf32, offset : 0, strides : [?, 1]>)
+    -> memref<?xf32> {
+  %1 = memref.expand_shape %arg0 [[0, 1], [2]] :
+    memref<?x?xf32, offset : 0, strides : [?, 1]> into
+    memref<?x4x?xf32, offset : 0, strides : [?, ?, 1]>
+  %2 = memref.collapse_shape %1 [[0, 1, 2]] :
+    memref<?x4x?xf32, offset : 0, strides : [?, ?, 1]> into
+    memref<?xf32>
+  return %2 : memref<?xf32>
+}
+// CHECK-LABEL: func @do_not_compose_collapse_of_expand_non_identity_layout
+// CHECK: expand
+// CHECK: collapse
+
+// -----
+
+func @compose_expand_of_expand(%arg0 : memref<?x?xf32>)
     -> memref<?x6x4x5x?xf32> {
   %0 = memref.expand_shape %arg0 [[0, 1], [2]]
       : memref<?x?xf32> into memref<?x4x?xf32>
   %1 = memref.expand_shape %0 [[0, 1], [2], [3, 4]]
       : memref<?x4x?xf32> into memref<?x6x4x5x?xf32>
   return %1 : memref<?x6x4x5x?xf32>
 }
-// CHECK-LABEL: expanding_memref_reshapes
+// CHECK-LABEL: func @compose_expand_of_expand
 //       CHECK:   memref.expand_shape %{{.*}} {{\[}}[0, 1, 2], [3, 4]]
 //   CHECK-NOT:   memref.expand_shape
 
 // -----
 
-func @expanding_memref_reshapes_to_zero_dim(%arg0 : memref<f32>)
-                                             -> memref<1x1x1xf32> {
+func @compose_expand_of_expand_of_zero_dim(%arg0 : memref<f32>)
   %0 = memref.expand_shape %arg0 [] : memref<f32> into memref<1xf32>
   %1 = memref.expand_shape %0 [[0, 1, 2]]
       : memref<1xf32> into memref<1x1x1xf32>
   return %1 : memref<1x1x1xf32>
 }
-// CHECK-LABEL: expanding_memref_reshapes_to_zero
+// CHECK-LABEL: func @compose_expand_of_expand_of_zero_dim
 //       CHECK:   memref.expand_shape %{{.*}} []
 //  CHECK-SAME:     memref<f32> into memref<1x1x1xf32>
 
 // -----
 
-func @fold_memref_reshape(%arg0 : memref<12x4xf32>) -> memref<12x4xf32> {
+func @fold_collapse_of_expand(%arg0 : memref<12x4xf32>) -> memref<12x4xf32> {
   %0 = memref.expand_shape %arg0 [[0, 1], [2]]
       : memref<12x4xf32> into memref<3x4x4xf32>
   %1 = memref.collapse_shape %0 [[0, 1], [2]]
       : memref<3x4x4xf32> into memref<12x4xf32>
   return %1 : memref<12x4xf32>
 }
-// CHECK-LABEL: @fold_memref_reshape
+// CHECK-LABEL: func @fold_collapse_of_expand
 //   CHECK-NOT:   linalg.{{.*}}_shape
 
 // -----
 
-func @fold_memref_reshape_dynamic(%arg0 : memref<?x?xf32>) -> memref<?x?xf32> {
+func @fold_collapse_collapse_of_expand(%arg0 : memref<?x?xf32>)
+    -> memref<?x?xf32> {
   %0 = memref.expand_shape %arg0 [[0, 1], [2]]
       : memref<?x?xf32> into memref<?x4x?xf32>
   %1 = memref.collapse_shape %0 [[0, 1], [2]]
       : memref<?x4x?xf32> into memref<?x?xf32>
   return %1 : memref<?x?xf32>
 }
-// CHECK-LABEL: @fold_memref_reshape_dynamic
+// CHECK-LABEL: @fold_collapse_collapse_of_expand
 //   CHECK-NOT:   linalg.{{.*}}_shape
 
 // -----
 
 func @fold_memref_expand_cast(%arg0 : memref<?x?xf32>) -> memref<2x4x4xf32> {
   %0 = memref.cast %arg0 : memref<?x?xf32> to memref<8x4xf32>
   %1 = memref.expand_shape %0 [[0, 1], [2]]
       : memref<8x4xf32> into memref<2x4x4xf32>

mlir/test/Dialect/Tensor/canonicalize.mlir

func @fold_overlapping_insert(%input : tensor<?x?x?xf32>, %slice1: tensor<4x?x8xf32>, %slice2: tensor<4x?x8xf32>, %i: index, %size: index) -> (tensor<?x?x?xf32>) {
   // CHECK: %[[INSERT:.+]] = tensor.insert_slice %[[SLICE2]] into %[[INPUT]]
   %1 = tensor.insert_slice %slice2 into %0[%c0, %i, 0] [4, %size, 8] [1, 1, %c1] : tensor<4x?x8xf32> into tensor<?x?x?xf32>
   // CHECK: return %[[INSERT]]
   return %1 : tensor<?x?x?xf32>
 }
 
 // -----
 
-func @expanding_tensor_reshapes(%arg0 : tensor<?x?xf32>)
+func @compose_expand_of_expand(%arg0 : tensor<?x?xf32>)
     -> tensor<?x6x4x?x5xf32> {
   %0 = tensor.expand_shape %arg0 [[0, 1], [2]]
       : tensor<?x?xf32> into tensor<?x4x?xf32>
   %1 = tensor.expand_shape %0 [[0, 1], [2], [3, 4]]
       : tensor<?x4x?xf32> into tensor<?x6x4x?x5xf32>
   return %1 : tensor<?x6x4x?x5xf32>
 }
-// CHECK-LABEL: expanding_tensor_reshapes
+// CHECK-LABEL: compose_expand_of_expand
 //       CHECK:   tensor.expand_shape %{{.*}} {{\[}}[0, 1, 2], [3, 4]]
 //   CHECK-NOT:   tensor.expand_shape
 
 // -----
 
-func @expanding_tensor_reshapes_to_zero_dim(%arg0 : tensor<f32>)
+func @compose_expand_of_expand_of_zero_dim(%arg0 : tensor<f32>)
     -> tensor<1x1x1xf32> {
   %0 = tensor.expand_shape %arg0 [] : tensor<f32> into tensor<1xf32>
   %1 = tensor.expand_shape %0 [[0, 1, 2]]
       : tensor<1xf32> into tensor<1x1x1xf32>
   return %1 : tensor<1x1x1xf32>
 }
-// CHECK-LABEL: expanding_tensor_reshapes_to_zero
+// CHECK-LABEL: compose_expand_of_expand_of_zero_dim
 //       CHECK:   tensor.expand_shape %{{.*}} []
 //  CHECK-SAME:     tensor<f32> into tensor<1x1x1xf32>
 
 // -----
 
-func @fold_tensor_reshape(%arg0 : tensor<12x4xf32>) -> tensor<12x4xf32> {
+func @fold_collapse_of_expand(%arg0 : tensor<12x4xf32>) -> tensor<12x4xf32> {
   %0 = tensor.expand_shape %arg0 [[0, 1], [2]]
       : tensor<12x4xf32> into tensor<3x4x4xf32>
   %1 = tensor.collapse_shape %0 [[0, 1], [2]]
       : tensor<3x4x4xf32> into tensor<12x4xf32>
   return %1 : tensor<12x4xf32>
 }
-// CHECK-LABEL: @fold_tensor_reshape
+// CHECK-LABEL: @fold_collapse_of_expand
 //   CHECK-NOT:   linalg.{{.*}}shape
 
 // -----
 
-func @fold_tensor_reshape_dynamic(%arg0 : tensor<?x?xf32>) -> tensor<?x?xf32> {
+func @fold_collapse_of_expand_dynamic(%arg0 : tensor<?x?xf32>)
+    -> tensor<?x?xf32> {
   %0 = tensor.expand_shape %arg0 [[0, 1], [2]]
       : tensor<?x?xf32> into tensor<?x4x?xf32>
   %1 = tensor.collapse_shape %0 [[0, 1], [2]]
       : tensor<?x4x?xf32> into tensor<?x?xf32>
   return %1 : tensor<?x?xf32>
 }
-// CHECK-LABEL: @fold_tensor_reshape_dynamic
+// CHECK-LABEL: @fold_collapse_of_expand_dynamic
 //   CHECK-NOT:   linalg.{{.*}}_shape
 
 // -----
-func @reshape_collapse(%arg0 : tensor<2x3x4x5x6x7x8xf32>)
+
+func @compose_expand_of_collapse(%arg0 : tensor<2x3x4x5x6x7x8xf32>)
     -> tensor<24x5x42x8xf32> {
   %0 = tensor.collapse_shape %arg0 [[0, 1, 2, 3, 4, 5, 6]]
       : tensor<2x3x4x5x6x7x8xf32> into tensor<40320xf32>
   %1 = tensor.expand_shape %0 [[0, 1, 2, 3]]
       : tensor<40320xf32> into tensor<24x5x42x8xf32>
   return %1 : tensor<24x5x42x8xf32>
 }
-//      CHECK: func @reshape_collapse
+//      CHECK: func @compose_expand_of_collapse
 // CHECK-SAME:   %[[ARG0:.+]]: tensor<2x3x4x5x6x7x8xf32>
 //      CHECK:   %[[RESULT:.+]] = tensor.collapse_shape %[[ARG0]]
 // CHECK-SAME:     [0, 1, 2], [3], [4, 5], [6]
 //      CHECK:   return %[[RESULT]]
 
 // -----
 
-func @reshape_expand(%arg0 : tensor<24x5x42x8xf32>)
+func @compose_expand_of_collapse_7D(%arg0 : tensor<24x5x42x8xf32>)
     -> tensor<2x3x4x5x6x7x8xf32> {
   %0 = tensor.collapse_shape %arg0 [[0, 1, 2, 3]]
       : tensor<24x5x42x8xf32> into tensor<40320xf32>
   %1 = tensor.expand_shape %0 [[0, 1, 2, 3, 4, 5, 6]]
       : tensor<40320xf32> into tensor<2x3x4x5x6x7x8xf32>
   return %1 : tensor<2x3x4x5x6x7x8xf32>
 }
-//      CHECK: func @reshape_expand
+//      CHECK: func @compose_expand_of_collapse_7D
 // CHECK-SAME:   %[[ARG0:.+]]: tensor<24x5x42x8xf32>
 //      CHECK:   %[[RESULT:.+]] = tensor.expand_shape %[[ARG0]]
 // CHECK-SAME:     [0, 1, 2], [3], [4, 5], [6]
 //      CHECK:   return %[[RESULT]]
 
 // -----
 
-func @expand_reshape_1D(%arg0 : tensor<2048xf32>) -> tensor<4x512xf32> {
+func @compose_collapse_of_expand(%arg : tensor<?x?x?xi64>)
+    -> tensor<?x?xi64> {
+  %0 = tensor.expand_shape %arg [[0], [1], [2, 3]]
+    : tensor<?x?x?xi64> into tensor<?x?x?x1xi64>
+  %1 = tensor.collapse_shape %0 [[0, 1], [2, 3]]
+    : tensor<?x?x?x1xi64> into tensor<?x?xi64>
+  return %1 : tensor<?x?xi64>
+}
+// CHECK-LABEL: func @compose_collapse_of_expand
+//       CHECK:   (%[[ARG:.*]]: tensor<?x?x?xi64>)
+//  CHECK-NEXT: tensor.collapse_shape %[[ARG]]
+//  CHECK-SAME:   [0, 1], [2]
+//  CHECK-SAME:   : tensor<?x?x?xi64> into tensor<?x?xi64>
+
+// -----
+
+func @compose_collapse_of_expand_1D(%arg0 : tensor<2048xf32>)
+    -> tensor<4x512xf32> {
   %0 = tensor.expand_shape %arg0 [[0, 1, 2, 3]]
     : tensor<2048xf32> into tensor<1x4x1x512xf32>
   %1 = tensor.collapse_shape %0 [[0, 1, 2], [3]]
     : tensor<1x4x1x512xf32> into tensor<4x512xf32>
   return %1 : tensor<4x512xf32>
 }
-//       CHECK: func @expand_reshape_1D
+//       CHECK: func @compose_collapse_of_expand_1D
 //       CHECK: tensor.expand_shape %{{.*}} {{\[}}[0, 1]]
 //  CHECK-SAME:   tensor<2048xf32> into tensor<4x512xf32>
 
 // -----
 
-// CHECK-LABEL: zero_rank_reshape_multi
+// CHECK-LABEL: func @zero_rank_reshape_multi
 func @zero_rank_reshape_multi(%arg0: tensor<f32>) -> tensor<f32> {
   // CHECK: return %arg0
   %0 = tensor.expand_shape %arg0 [] : tensor<f32> into tensor<1xf32>
   %1 = tensor.expand_shape %0 [[0, 1]] : tensor<1xf32> into tensor<1x1xf32>
   %2 = tensor.collapse_shape %1 [] : tensor<1x1xf32> into tensor<f32>
   return %2 : tensor<f32>
 }
 
 // -----
 
-func @collapsing_tensor_reshapes(%arg0 : tensor<?x?x?x?x?xf32>)
+func @compose_collapse_of_collapse(%arg0 : tensor<?x?x?x?x?xf32>)
     -> tensor<?x?xf32> {
   %0 = tensor.collapse_shape %arg0 [[0, 1], [2], [3, 4]]
       : tensor<?x?x?x?x?xf32> into tensor<?x?x?xf32>
   %1 = tensor.collapse_shape %0 [[0, 1], [2]]
       : tensor<?x?x?xf32> into tensor<?x?xf32>
   return %1 : tensor<?x?xf32>
 }
-// CHECK-LABEL: collapsing_tensor_reshapes
+// CHECK-LABEL: func @compose_collapse_of_collapse
 //       CHECK:   tensor.collapse_shape %{{.*}} {{\[}}[0, 1, 2], [3, 4]]
 //   CHECK-NOT:   tensor.collapse_shape
 
 // -----
 
-func @collapsing_tensor_reshapes_to_zero_dim(%arg0 : tensor<1x1x1xf32>)
+func @compose_collapse_of_collapse_zero_dim(%arg0 : tensor<1x1x1xf32>)
     -> tensor<f32> {
   %0 = tensor.collapse_shape %arg0 [[0, 1, 2]]
       : tensor<1x1x1xf32> into tensor<1xf32>
   %1 = tensor.collapse_shape %0 [] : tensor<1xf32> into tensor<f32>
   return %1 : tensor<f32>
 }
-// CHECK-LABEL: collapsing_tensor_reshapes_to_zero
+// CHECK-LABEL: func @compose_collapse_of_collapse_zero_dim
 //       CHECK:   tensor.collapse_shape %{{.*}} []
 //  CHECK-SAME:     tensor<1x1x1xf32> into tensor<f32>
 
 // -----
 
-func @fold_reshape_1D(%arg0 : tensor<4x512xf32>) -> tensor<2048xf32> {
+func @fold_collapse_of_expand_1D(%arg0 : tensor<4x512xf32>) -> tensor<2048xf32> {
   %0 = tensor.expand_shape %arg0 [[0, 1, 2], [3]]
     : tensor<4x512xf32> into tensor<1x4x1x512xf32>
   %1 = tensor.collapse_shape %0 [[0, 1, 2, 3]]
     : tensor<1x4x1x512xf32> into tensor<2048xf32>
   return %1 : tensor<2048xf32>
 }
-//       CHECK: func @fold_reshape_1D
+//       CHECK: func @fold_collapse_of_expand_1D
 //       CHECK: tensor.collapse_shape %{{.*}} {{\[}}[0, 1]]
 //  CHECK-SAME:   tensor<4x512xf32> into tensor<2048xf32>
 
 // -----
 
-func @fold_reshape_unit_dims(%arg0 : tensor<2048x1x1xf32>)
+func @fold_collapse_of_expand_unit_dims(%arg0 : tensor<2048x1x1xf32>)
     -> tensor<4x512x1x1xf32> {
   %0 = tensor.expand_shape %arg0 [[0, 1, 2, 3], [4], [5]]
     : tensor<2048x1x1xf32> into tensor<1x4x1x512x1x1xf32>
   %1 = tensor.collapse_shape %0 [[0, 1, 2], [3], [4], [5]]
     : tensor<1x4x1x512x1x1xf32> into tensor<4x512x1x1xf32>
   return %1 : tensor<4x512x1x1xf32>
 }
-//       CHECK: func @fold_reshape_unit_dims
+//       CHECK: func @fold_collapse_of_expand_unit_dims
 //       CHECK: tensor.expand_shape %{{.*}} {{\[}}[0, 1], [2], [3]]
 //  CHECK-SAME:   tensor<2048x1x1xf32> into tensor<4x512x1x1xf32>
 
 // -----
 
-func @expand_reshape_unit_dims(%arg0 : tensor<2048x1x2048xf32>)
+func @compose_collapse_of_expand_unit_dims(%arg0 : tensor<2048x1x2048xf32>)
     -> tensor<4x512x1x512x4xf32> {
   %0 = tensor.expand_shape %arg0 [[0, 1, 2, 3, 4], [5], [6, 7, 8]]
     : tensor<2048x1x2048xf32> into tensor<1x4x1x512x1x1x512x1x4xf32>
   %1 = tensor.collapse_shape %0 [[0, 1, 2], [3, 4], [5], [6, 7], [8]]
     : tensor<1x4x1x512x1x1x512x1x4xf32> into tensor<4x512x1x512x4xf32>
   return %1 : tensor<4x512x1x512x4xf32>
 }
-//       CHECK: func @expand_reshape_unit_dims
+//       CHECK: func @compose_collapse_of_expand_unit_dims
 //       CHECK: tensor.expand_shape %{{.*}} {{\[}}[0, 1], [2], [3, 4]]
 //  CHECK-SAME:   tensor<2048x1x2048xf32> into tensor<4x512x1x512x4xf32>
 
 // -----
 
-func @fold_reshape_trailing_unit_dims(%arg0: tensor<2xf32>) -> tensor<2x1xf32> {
+func @compose_collapse_of_expand_trailing_unit_dims(%arg0: tensor<2xf32>)
+    -> tensor<2x1xf32> {
   %0 = tensor.expand_shape %arg0 [[0, 1, 2]]
       : tensor<2xf32> into tensor<2x1x1xf32>
   %1 = tensor.collapse_shape %0 [[0], [1, 2]]
       : tensor<2x1x1xf32> into tensor<2x1xf32>
   return %1 : tensor<2x1xf32>
 }
-//       CHECK: func @fold_reshape_trailing_unit_dims
+//       CHECK: func @compose_collapse_of_expand_trailing_unit_dims
 //       CHECK: tensor.expand_shape %{{.*}} {{\[}}[0, 1]]
 //  CHECK-SAME:   tensor<2xf32> into tensor<2x1xf32>
 
 // -----
 
-func @collapse_reshape_unit_dims_dynamic(%arg0 : tensor<?x1x?x1x1x?x?x1x1xf32>)
-    -> tensor<?x?x?x?xf32> {
+func @compose_collapse_of_collapse_unit_dims_dynamic(
+    %arg0 : tensor<?x1x?x1x1x?x?x1x1xf32>) -> tensor<?x?x?x?xf32> {
   %0 = tensor.collapse_shape %arg0 [[0], [1, 2], [3], [4], [5], [6, 7, 8]]
     : tensor<?x1x?x1x1x?x?x1x1xf32> into tensor<?x?x1x1x?x?xf32>
   %1 = tensor.collapse_shape %0 [[0], [1], [2, 3, 4], [5]]
     : tensor<?x?x1x1x?x?xf32> into tensor<?x?x?x?xf32>
   return %1 : tensor<?x?x?x?xf32>
 }
-//       CHECK: func @collapse_reshape_unit_dims_dynamic
+//       CHECK: func @compose_collapse_of_collapse_unit_dims_dynamic
 //       CHECK: tensor.collapse_shape
 //  CHECK-SAME:   [0], [1, 2], [3, 4, 5], [6, 7, 8]
 //  CHECK-SAME:   tensor<?x1x?x1x1x?x?x1x1xf32> into tensor<?x?x?x?xf32>
 
 // -----
 
-func @fold_reshape_trailing_unit_dims(%arg0: tensor<2xf32>) -> tensor<2x1xf32>
-{
+func @fold_collapse_of_expand_trailing_unit_dims(%arg0: tensor<2xf32>)
+    -> tensor<2x1xf32> {
   %0 = tensor.expand_shape %arg0 [[0, 1, 2]]
       : tensor<2xf32> into tensor<2x1x1xf32>
   %1 = tensor.collapse_shape %0 [[0], [1, 2]]
       : tensor<2x1x1xf32> into tensor<2x1xf32>
   return %1 : tensor<2x1xf32>
 }
-//       CHECK: func @fold_reshape_trailing_unit_dims
+//       CHECK: func @fold_collapse_of_expand_trailing_unit_dims
 //       CHECK: tensor.expand_shape %{{.*}} {{\[}}[0, 1]]
 //  CHECK-SAME:   tensor<2xf32> into tensor<2x1xf32>
 
 // -----
 
-func @fold_reshape_trailing_unit_dims_dynamic(%arg0: tensor<1x1x?x1x1x1xf32>)
-    -> tensor<?xf32> {
+func @fold_collapse_of_collapse_trailing_unit_dims_dynamic(
+    %arg0: tensor<1x1x?x1x1x1xf32>) -> tensor<?xf32> {
   %0 = tensor.collapse_shape %arg0 [[0, 1, 2], [3], [4], [5]]
       : tensor<1x1x?x1x1x1xf32> into tensor<?x1x1x1xf32>
   %1 = tensor.collapse_shape %0 [[0, 1, 2, 3]]
       : tensor<?x1x1x1xf32> into tensor<?xf32>
   return %1 : tensor<?xf32>
 }
-//       CHECK: func @fold_reshape_trailing_unit_dims_dynamic
+//       CHECK: func @fold_collapse_of_collapse_trailing_unit_dims_dynamic
 //       CHECK: 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>)
+func @fold_collapse_of_expand_trailing_unit_dims(%arg0: tensor<12x42x1x1xf32>)
     -> tensor<12x42xf32> {
   %0 = tensor.expand_shape %arg0 [[0], [1], [2], [3, 4]]
       : tensor<12x42x1x1xf32> into tensor<12x42x1x1x1xf32>
   %1 = 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: func @fold_collapse_of_expand_trailing_unit_dims
 //       CHECK: tensor.collapse_shape %{{.*}} {{\[}}[0], [1, 2, 3]]
 //  CHECK-SAME:   tensor<12x42x1x1xf32> into tensor<12x42xf32>
 
 // -----
 
-func @fold_reshapes_unit_dims_in_middle(%arg0 : tensor<?x?x?xf32>) -> tensor<?x?xf32> {
+func @fold_collapse_of_expand_unit_dims_in_middle(%arg0 : tensor<?x?x?xf32>)
+    -> tensor<?x?xf32> {
   %0 = tensor.expand_shape %arg0 [[0], [1], [2, 3]]
       : tensor<?x?x?xf32> into tensor<?x?x1x?xf32>
   %1 = tensor.collapse_shape %0 [[0], [1, 2, 3]]
       : tensor<?x?x1x?xf32> into tensor<?x?xf32>
   return %1 : tensor<?x?xf32>
 }
-// CHECK-LABEL: func @fold_reshapes_unit_dims_in_middle
+// CHECK-LABEL: func @fold_collapse_of_expand_unit_dims_in_middle
 //  CHECK-SAME: (%[[ARG:.*]]: tensor<?x?x?xf32>
 //       CHECK: tensor.collapse_shape %[[ARG]] {{\[}}[0], [1, 2]]
 //  CHECK-SAME:   tensor<?x?x?xf32> into tensor<?x?xf32>
 
 // -----
 
-func @no_fold_reshape_incompatible(%arg0 : tensor<4x6x8xf32>)
+func @no_fold_collapse_of_expand_incompatible(%arg0 : tensor<4x6x8xf32>)
     -> tensor<2x6x16xf32> {
   %0 = tensor.expand_shape %arg0 [[0, 1], [2, 3], [4]]
       : tensor<4x6x8xf32> into tensor<2x2x3x2x8xf32>
   %1 = tensor.collapse_shape %0 [[0], [1, 2], [3, 4]]
       : tensor<2x2x3x2x8xf32> into tensor<2x6x16xf32>
   return %1 : tensor<2x6x16xf32>
 }
-// CHECK-LABEL: func @no_fold_reshape_incompatible
+// CHECK-LABEL: func @no_fold_collapse_of_expand_incompatible
 //       CHECK:   tensor.expand_shape
 //       CHECK:   tensor.collapse_shape
 
 // -----
 
-func @no_fold_reshape_empty_expr(%arg0: tensor<3x2x2xf32>) -> tensor<12x1xf32> {
+func @no_fold_collapse_of_expand_empty_expr(%arg0: tensor<3x2x2xf32>)
+    -> tensor<12x1xf32> {
   %0 = tensor.expand_shape %arg0 [[0], [1], [2, 3]]
       : tensor<3x2x2xf32> into tensor<3x2x2x1xf32>
   %1 = tensor.collapse_shape %0 [[0, 1, 2], [3]]
       : tensor<3x2x2x1xf32> into tensor<12x1xf32>
   return %1 : tensor<12x1xf32>
 }
-//      CHECK: func @no_fold_reshape_empty_expr
+//      CHECK: func @no_fold_collapse_of_expand_empty_expr
 // CHECK-SAME:    %[[ARG0:.+]]: tensor<3x2x2xf32>
 //      CHECK:    %[[RARG0:.+]] = tensor.expand_shape %[[ARG0]]
 // CHECK-SAME:      [0], [1], [2, 3]
 //      CHECK:    %[[RES:.+]] = tensor.collapse_shape %[[RARG0]]
 // CHECK-SAME:      [0, 1, 2], [3]
 //      CHECK:    return %[[RES:.+]] : tensor<12x1xf32>
 
 // -----
func @fold_rank() -> (index) {
   %rank_0 = tensor.rank %const_0 : tensor<2x1x4xi32>
 
   // CHECK-NEXT: return [[C3]]
   return %rank_0 : index
 }
 
 // -----
 
-// CHECK-LABEL: func @pad_tensor_same_static_shape(
+// CHECK-LABEL: func @pad_same_static_shape(
 //  CHECK-SAME:   %[[ARG0:.*]]: tensor<5x6xf32>
 //   CHECK-NOT:   tensor.pad
 //       CHECK:   return %[[ARG0]]
-func @pad_tensor_same_static_shape(%arg0: tensor<5x6xf32>, %a: index)
+func @pad_same_static_shape(%arg0: tensor<5x6xf32>, %a: index)
     -> tensor<5x6xf32> {
   %cst = arith.constant 0.000000e+00 : f32
   %0 = tensor.pad %arg0 low[%a, 0] high[0, %a] {
         ^bb0(%arg1: index, %arg2: index):
           tensor.yield %cst : f32
   } : tensor<5x6xf32> to tensor<5x6xf32>
   return %0 : tensor<5x6xf32>
 }
 
 // -----
 
-// CHECK-LABEL: func @pad_tensor_nofold_same_static_shape(
+// CHECK-LABEL: func @pad_nofold_same_static_shape(
 //  CHECK-SAME:   %[[ARG0:.*]]: tensor<5x6xf32>
 //       CHECK:   %[[PAD:.*]] = tensor.pad
 //       CHECK:   return %[[PAD]]
-func @pad_tensor_nofold_same_static_shape(%arg0: tensor<5x6xf32>, %a: index)
+func @pad_nofold_same_static_shape(%arg0: tensor<5x6xf32>, %a: index)
     -> tensor<5x6xf32> {
   %cst = arith.constant 0.000000e+00 : f32
   %0 = tensor.pad %arg0 nofold low[%a, 0] high[0, %a] {
         ^bb0(%arg1: index, %arg2: index):
           tensor.yield %cst : f32
   } : tensor<5x6xf32> to tensor<5x6xf32>
   return %0 : tensor<5x6xf32>
 }
 
 // -----
 
-// CHECK-LABEL:   func @pad_tensor_after_cast_different_shape(
+// CHECK-LABEL:   func @pad_after_cast_different_shape(
 // CHECK-SAME:      %[[INPUT:.*]]: tensor<?x64x?x?xf32>) -> tensor<?x?x?x?xf32> {
 // CHECK:           %[[CST:.*]] = arith.constant 0.000000e+00 : f32
 // CHECK:           %[[PADDED:.*]] = tensor.pad %[[INPUT]]
 // CHECK-SAME:        low[0, 0, 1, 1] high[0, 0, 1, 1]  {
 // CHECK:           ^bb0(%[[ARG1:.*]]: index, %[[ARG2:.*]]: index, %[[ARG3:.*]]: index, %[[ARG4:.*]]: index):
 // CHECK:             tensor.yield %[[CST]] : f32
 // CHECK:           } : tensor<?x64x?x?xf32> to tensor<?x64x?x?xf32>
 // CHECK:           %[[DYNAMIC:.*]] = tensor.cast %[[PADDED:.*]] :
 // CHECK-SAME:         tensor<?x64x?x?xf32> to tensor<?x?x?x?xf32>
 // CHECK:           return %[[DYNAMIC]] : tensor<?x?x?x?xf32>
 // CHECK:         }
-func @pad_tensor_after_cast_different_shape(%arg0: tensor<?x64x?x?xf32>)
+func @pad_after_cast_different_shape(%arg0: tensor<?x64x?x?xf32>)
     -> tensor<?x?x?x?xf32> {
   %cst = arith.constant 0.000000e+00 : f32
   %dynamic = tensor.cast %arg0 : tensor<?x64x?x?xf32> to tensor<?x?x?x?xf32>
   %padded = tensor.pad %dynamic low[0, 0, 1, 1] high[0, 0, 1, 1]  {
     ^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):
     tensor.yield %cst: f32
   } : tensor<?x?x?x?xf32> to tensor<?x?x?x?xf32>
   return %padded: tensor<?x?x?x?xf32>
 }
 
 // -----
 
-// CHECK-LABEL:   func @pad_tensor_after_cast_same_shape(
+// CHECK-LABEL:   func @pad_after_cast_same_shape(
 // CHECK-SAME:      %[[INPUT:.*]]: tensor<?x64x?x?xf32>,
 // CHECK-SAME:      %[[PADDING:.*]]: index) -> tensor<?x?x?x?xf32> {
 // CHECK:           %[[CST:.*]] = arith.constant 0.000000e+00 : f32
 // CHECK:           %[[PADDED:.*]] = tensor.pad %[[INPUT]]
 // CHECK-SAME:        low[0, %[[PADDING]], 1, 1] high[0, %[[PADDING]], 1, 1]  {
 // CHECK:           ^bb0(%[[ARG1:.*]]: index, %[[ARG2:.*]]: index, %[[ARG3:.*]]: index, %[[ARG4:.*]]: index):
 // CHECK:             tensor.yield %[[CST]] : f32
 // CHECK:           } : tensor<?x64x?x?xf32> to tensor<?x?x?x?xf32>
 // CHECK:           return %[[PADDED:.*]] : tensor<?x?x?x?xf32>
 // CHECK:         }
-func @pad_tensor_after_cast_same_shape(%arg0: tensor<?x64x?x?xf32>, %padding : index)
+func @pad_after_cast_same_shape(%arg0: tensor<?x64x?x?xf32>, %padding : index)
     -> tensor<?x?x?x?xf32> {
   %cst = arith.constant 0.000000e+00 : f32
   %dynamic = tensor.cast %arg0 : tensor<?x64x?x?xf32> to tensor<?x?x?x?xf32>
   %padded = tensor.pad %dynamic low[0, %padding, 1, 1] high[0, %padding, 1, 1]  {
     ^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):
     tensor.yield %cst: f32
   } : tensor<?x?x?x?xf32> to tensor<?x?x?x?xf32>
   return %padded: tensor<?x?x?x?xf32>
 }
 
 // -----
 
-// CHECK-LABEL: func @pad_tensor_of_cast(
+// CHECK-LABEL: func @pad_of_cast(
 // CHECK-NOT:     tensor.cast
 // CHECK:         tensor.pad
 // CHECK:         tensor<8x?xf32> to tensor<8x32xf32>
-func @pad_tensor_of_cast(%t: tensor<8x?xf32>, %s: index) -> tensor<8x32xf32> {
+func @pad_of_cast(%t: tensor<8x?xf32>, %s: index) -> tensor<8x32xf32> {
   %c0 = arith.constant 0 : index
   %cst = arith.constant 0.000000e+00 : f32
   %0 = tensor.cast %t : tensor<8x?xf32> to tensor<?x?xf32>
   %1 = tensor.pad %0 low[%c0, %c0] high[%c0, %s]  {
   ^bb0(%arg9: index, %arg10: index):
     tensor.yield %cst : f32
   } : tensor<?x?xf32> to tensor<8x32xf32>
   return %1 : tensor<8x32xf32>
func @cast_of_pad_less_static(%arg0: tensor<32x?x?xf32>, %padding: index) -> tensor<?x32x32xf32> {
   // CHECK: %[[CAST:.*]] = tensor.cast
   %casted = tensor.cast %padded : tensor<32x?x?xf32> to tensor<?x32x32xf32>
   // CHECK: return %[[CAST]]
   return %casted : tensor<?x32x32xf32>
 }
 
 // -----
 
-func @tensor_pad_cast_fold(%arg0: tensor<4x4xf32>) -> tensor<4x4xf32> {
+func @pad_cast_fold(%arg0: tensor<4x4xf32>) -> tensor<4x4xf32> {
   %c0 = arith.constant 0 : index
   %cst = arith.constant 0.0 : f32
   %0 = tensor.cast %arg0 : tensor<4x4xf32> to tensor<?x?xf32>
   %1 = tensor.pad %0 low[%c0, %c0] high[%c0, %c0]  {
     ^bb0(%arg1: index, %arg2: index):
       tensor.yield %cst : f32
   } : tensor<?x?xf32> to tensor<4x4xf32>
   return %1 : tensor<4x4xf32>
 }
-// CHECK-LABEL: @tensor_pad_cast
+// CHECK-LABEL: @pad_cast
 // CHECK-SAME: %[[ARG0:.+]]: tensor<4x4xf32>
 // CHECK: return %[[ARG0]]
 
 // -----
 
-// CHECK-LABEL: func @fold_pad_tensor_source_cast(
+// CHECK-LABEL: func @fold_pad_source_cast(
 //  CHECK-SAME:                  %[[ARG0:.*]]: tensor<4x?xf32>
 //   CHECK-NOT:   tensor.cast
 //       CHECK:   %[[RESULT:.*]] = tensor.pad %[[ARG0]]
-func @fold_pad_tensor_source_cast(%arg0: tensor<4x?xf32>) -> tensor<4x4xf32> {
+func @fold_pad_source_cast(%arg0: tensor<4x?xf32>) -> tensor<4x4xf32> {
   %cst = arith.constant 0.0 : f32
   %0 = tensor.cast %arg0 : tensor<4x?xf32> to tensor<?x?xf32>
   %1 = tensor.pad %0 low[0, 0] high[0, 1]  {
     ^bb0(%arg1: index, %arg2: index):
       tensor.yield %cst : f32
   } : tensor<?x?xf32> to tensor<4x4xf32>
   return %1 : tensor<4x4xf32>
 }