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

[hanchung]

Summary

This reverts commit 64f659bee67b5a024defeb3cd2ecf65e1ad8c0a7.

An invalid tensor.expand_shape op is generated 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>
}

[pifon2a]

Thank you for the IR to reproduce it. The 0-rank case was not handled correctly. And it is also surprising that there were no tests for that. I fixed it in https://github.com/llvm/llvm-project/commit/747b10be95200a71f4e6b8ca9f0aaea20db1b164

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 ComposeReassociativeReshapeOps : public OpRewritePattern<ReshapeOpTy> {
+struct CollapseReshapeOps : 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 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,
+/// 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,
                                 PatternRewriter &rewriter) const override {
-    auto expandOp = collapseOp.src().template getDefiningOp<ExpandOpTy>();
-    if (!expandOp)
-      return failure();
-
-    ShapedType srcType = expandOp.getSrcType();
-    ShapedType resultType = collapseOp.getResultType();
-
-    if (hasNonIdentityLayout(collapseOp.src().getType()) ||
-        hasNonIdentityLayout(expandOp.src().getType()) ||
-        hasNonIdentityLayout(expandOp.result().getType()))
-      return failure();
-
-    int64_t srcRank = srcType.getRank();
-    int64_t resultRank = resultType.getRank();
-    if (srcType == resultType)
+    auto srcReshapeOp =
+        reshapeOp.src().template getDefiningOp<InverseReshapeOpTy>();
+    if (!srcReshapeOp)
       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();
-    }
+    ShapedType srcReshapeSrcType = srcReshapeOp.getSrcType();
+    ShapedType intermediateType = reshapeOp.getSrcType();
+    ShapedType resultType = reshapeOp.getResultType();
 
-    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);
+    // 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)
+    bool originalOpExpands =
+        intermediateType.getRank() > srcReshapeSrcType.getRank();
+    bool resultingOpExpands =
+        resultType.getRank() > srcReshapeSrcType.getRank();
+    if (!(resultingOpExpands ^ originalOpExpands))
+      rewriter.replaceOpWithNewOp<InverseReshapeOpTy>(
+          reshapeOp, resultType, srcReshapeOp.src(), *reassociationIndices);
     else
-      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);
+      rewriter.replaceOpWithNewOp<ReshapeOpTy>(
+          reshapeOp, resultType, srcReshapeOp.src(), *reassociationIndices);
     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<ComposeReassociativeReshapeOps<ExpandShapeOp>,
-              ComposeExpandOfCollapseOp<ExpandShapeOp, CollapseShapeOp>>(
-      context);
+  results.add<CollapseReshapeOps<ExpandShapeOp>,
+              CollapseMixedReshapeOps<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 (newResultType == op.getResultType()) {
       rewriter.replaceOpWithNewOp<CastOp>(op, op.getType(), newOp);
     }
     return success();
   }
 };
 
 void CollapseShapeOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                                   MLIRContext *context) {
-  results.add<ComposeReassociativeReshapeOps<CollapseShapeOp>,
-              ComposeCollapseOfExpandOp<CollapseShapeOp, ExpandShapeOp>,
+  results.add<CollapseReshapeOps<CollapseShapeOp>,
+              CollapseMixedReshapeOps<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<ComposeReassociativeReshapeOps<ExpandShapeOp>,
-              ComposeExpandOfCollapseOp<ExpandShapeOp, CollapseShapeOp>,
+  results.add<CollapseReshapeOps<ExpandShapeOp>,
+              CollapseMixedReshapeOps<ExpandShapeOp, CollapseShapeOp>,
               FoldReshapeWithConstant<ExpandShapeOp>,
               FoldReshapeWithFromElements<ExpandShapeOp>>(context);
 }
 
 void CollapseShapeOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                                   MLIRContext *context) {
-  results.add<ComposeReassociativeReshapeOps<CollapseShapeOp>,
-              ComposeCollapseOfExpandOp<CollapseShapeOp, ExpandShapeOp>,
+  results.add<CollapseReshapeOps<CollapseShapeOp>,
+              CollapseMixedReshapeOps<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) {
-  if (sourceType.getRank() > targetType.getRank())
-    return getReassociationIndicesForCollapse(sourceType.getShape(),
-                                              targetType.getShape());
-  if (sourceType.getRank() < targetType.getRank())
-    return getReassociationIndicesForCollapse(targetType.getShape(),
-                                              sourceType.getShape());
+  // 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())
+    std::swap(sourceType, targetType);
 
-Optional<SmallVector<ReassociationIndices>>
-mlir::getReassociationIndicesForCollapse(ArrayRef<int64_t> sourceShape,
-                                         ArrayRef<int64_t> targetShape) {
-  if (sourceShape.size() <= targetShape.size())
-    return llvm::None;
+  ArrayRef<int64_t> sourceShape = sourceType.getShape();
+  ArrayRef<int64_t> targetShape = targetType.getShape();
   unsigned sourceDim = 0;
   SmallVector<ReassociationIndices> reassociationMap;
-  reassociationMap.reserve(targetShape.size());
+  reassociationMap.reserve(targetType.getRank());
 
   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 == targetShape.size())
+    if (targetDim == targetType.getRank())
       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 @compose_collapse_of_collapse_zero_dim(%arg0 : memref<1x1x1xf32>)
+func @collapsing_memref_reshapes_to_zero_dim(%arg0 : memref<1x1x1xf32>)
                                              -> memref<f32> {
   %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: func @compose_collapse_of_collapse_zero_dim
+// CHECK-LABEL: collapsing_memref_reshapes_to_zero
 //       CHECK:   memref.collapse_shape %{{.*}} []
 //  CHECK-SAME:     memref<1x1x1xf32> into memref<f32>
 
 // -----
 
-func @compose_collapse_of_collapse(%arg0 : memref<?x?x?x?x?xf32>)
+func @collapsing_memref_reshapes(%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: func @compose_collapse_of_collapse
+// CHECK-LABEL: collapsing_memref_reshapes
 //       CHECK:   memref.collapse_shape %{{.*}} {{\[}}[0, 1, 2], [3, 4]]
 //   CHECK-NOT:   memref.collapse_shape
 
 // -----
 
-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>)
+func @expanding_memref_reshapes(%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: func @compose_expand_of_expand
+// CHECK-LABEL: expanding_memref_reshapes
 //       CHECK:   memref.expand_shape %{{.*}} {{\[}}[0, 1, 2], [3, 4]]
 //   CHECK-NOT:   memref.expand_shape
 
 // -----
 
-func @compose_expand_of_expand_of_zero_dim(%arg0 : memref<f32>)
+func @expanding_memref_reshapes_to_zero_dim(%arg0 : memref<f32>)
                                              -> memref<1x1x1xf32> {
   %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: func @compose_expand_of_expand_of_zero_dim
+// CHECK-LABEL: expanding_memref_reshapes_to_zero
 //       CHECK:   memref.expand_shape %{{.*}} []
 //  CHECK-SAME:     memref<f32> into memref<1x1x1xf32>
 
 // -----
 
-func @fold_collapse_of_expand(%arg0 : memref<12x4xf32>) -> memref<12x4xf32> {
+func @fold_memref_reshape(%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: func @fold_collapse_of_expand
+// CHECK-LABEL: @fold_memref_reshape
 //   CHECK-NOT:   linalg.{{.*}}_shape
 
 // -----
 
-func @fold_collapse_collapse_of_expand(%arg0 : memref<?x?xf32>)
-    -> memref<?x?xf32> {
+func @fold_memref_reshape_dynamic(%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_collapse_collapse_of_expand
+// CHECK-LABEL: @fold_memref_reshape_dynamic
 //   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 @compose_expand_of_expand(%arg0 : tensor<?x?xf32>)
+func @expanding_tensor_reshapes(%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: compose_expand_of_expand
+// CHECK-LABEL: expanding_tensor_reshapes
 //       CHECK:   tensor.expand_shape %{{.*}} {{\[}}[0, 1, 2], [3, 4]]
 //   CHECK-NOT:   tensor.expand_shape
 
 // -----
 
-func @compose_expand_of_expand_of_zero_dim(%arg0 : tensor<f32>)
+func @expanding_tensor_reshapes_to_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: compose_expand_of_expand_of_zero_dim
+// CHECK-LABEL: expanding_tensor_reshapes_to_zero
 //       CHECK:   tensor.expand_shape %{{.*}} []
 //  CHECK-SAME:     tensor<f32> into tensor<1x1x1xf32>
 
 // -----
 
-func @fold_collapse_of_expand(%arg0 : tensor<12x4xf32>) -> tensor<12x4xf32> {
+func @fold_tensor_reshape(%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_collapse_of_expand
+// CHECK-LABEL: @fold_tensor_reshape
 //   CHECK-NOT:   linalg.{{.*}}shape
 
 // -----
 
-func @fold_collapse_of_expand_dynamic(%arg0 : tensor<?x?xf32>)
-    -> tensor<?x?xf32> {
+func @fold_tensor_reshape_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_collapse_of_expand_dynamic
+// CHECK-LABEL: @fold_tensor_reshape_dynamic
 //   CHECK-NOT:   linalg.{{.*}}_shape
 
 // -----
-
-func @compose_expand_of_collapse(%arg0 : tensor<2x3x4x5x6x7x8xf32>)
+func @reshape_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 @compose_expand_of_collapse
+//      CHECK: func @reshape_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 @compose_expand_of_collapse_7D(%arg0 : tensor<24x5x42x8xf32>)
+func @reshape_expand(%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 @compose_expand_of_collapse_7D
+//      CHECK: func @reshape_expand
 // CHECK-SAME:   %[[ARG0:.+]]: tensor<24x5x42x8xf32>
 //      CHECK:   %[[RESULT:.+]] = tensor.expand_shape %[[ARG0]]
 // CHECK-SAME:     [0, 1, 2], [3], [4, 5], [6]
 //      CHECK:   return %[[RESULT]]
 
 // -----
 
-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> {
+func @expand_reshape_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 @compose_collapse_of_expand_1D
+//       CHECK: func @expand_reshape_1D
 //       CHECK: tensor.expand_shape %{{.*}} {{\[}}[0, 1]]
 //  CHECK-SAME:   tensor<2048xf32> into tensor<4x512xf32>
 
 // -----
 
-// CHECK-LABEL: func @zero_rank_reshape_multi
+// CHECK-LABEL: 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 @compose_collapse_of_collapse(%arg0 : tensor<?x?x?x?x?xf32>)
+func @collapsing_tensor_reshapes(%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: func @compose_collapse_of_collapse
+// CHECK-LABEL: collapsing_tensor_reshapes
 //       CHECK:   tensor.collapse_shape %{{.*}} {{\[}}[0, 1, 2], [3, 4]]
 //   CHECK-NOT:   tensor.collapse_shape
 
 // -----
 
-func @compose_collapse_of_collapse_zero_dim(%arg0 : tensor<1x1x1xf32>)
+func @collapsing_tensor_reshapes_to_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: func @compose_collapse_of_collapse_zero_dim
+// CHECK-LABEL: collapsing_tensor_reshapes_to_zero
 //       CHECK:   tensor.collapse_shape %{{.*}} []
 //  CHECK-SAME:     tensor<1x1x1xf32> into tensor<f32>
 
 // -----
 
-func @fold_collapse_of_expand_1D(%arg0 : tensor<4x512xf32>) -> tensor<2048xf32> {
+func @fold_reshape_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_collapse_of_expand_1D
+//       CHECK: func @fold_reshape_1D
 //       CHECK: tensor.collapse_shape %{{.*}} {{\[}}[0, 1]]
 //  CHECK-SAME:   tensor<4x512xf32> into tensor<2048xf32>
 
 // -----
 
-func @fold_collapse_of_expand_unit_dims(%arg0 : tensor<2048x1x1xf32>)
+func @fold_reshape_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_collapse_of_expand_unit_dims
+//       CHECK: func @fold_reshape_unit_dims
 //       CHECK: tensor.expand_shape %{{.*}} {{\[}}[0, 1], [2], [3]]
 //  CHECK-SAME:   tensor<2048x1x1xf32> into tensor<4x512x1x1xf32>
 
 // -----
 
-func @compose_collapse_of_expand_unit_dims(%arg0 : tensor<2048x1x2048xf32>)
+func @expand_reshape_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 @compose_collapse_of_expand_unit_dims
+//       CHECK: func @expand_reshape_unit_dims
 //       CHECK: tensor.expand_shape %{{.*}} {{\[}}[0, 1], [2], [3, 4]]
 //  CHECK-SAME:   tensor<2048x1x2048xf32> into tensor<4x512x1x512x4xf32>
 
 // -----
 
-func @compose_collapse_of_expand_trailing_unit_dims(%arg0: tensor<2xf32>)
-    -> tensor<2x1xf32> {
+func @fold_reshape_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 @compose_collapse_of_expand_trailing_unit_dims
+//       CHECK: func @fold_reshape_trailing_unit_dims
 //       CHECK: tensor.expand_shape %{{.*}} {{\[}}[0, 1]]
 //  CHECK-SAME:   tensor<2xf32> into tensor<2x1xf32>
 
 // -----
 
-func @compose_collapse_of_collapse_unit_dims_dynamic(
-    %arg0 : tensor<?x1x?x1x1x?x?x1x1xf32>) -> tensor<?x?x?x?xf32> {
+func @collapse_reshape_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 @compose_collapse_of_collapse_unit_dims_dynamic
+//       CHECK: func @collapse_reshape_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_collapse_of_expand_trailing_unit_dims(%arg0: tensor<2xf32>)
-    -> tensor<2x1xf32> {
+func @fold_reshape_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_collapse_of_expand_trailing_unit_dims
+//       CHECK: func @fold_reshape_trailing_unit_dims
 //       CHECK: tensor.expand_shape %{{.*}} {{\[}}[0, 1]]
 //  CHECK-SAME:   tensor<2xf32> into tensor<2x1xf32>
 
 // -----
 
-func @fold_collapse_of_collapse_trailing_unit_dims_dynamic(
-    %arg0: tensor<1x1x?x1x1x1xf32>) -> tensor<?xf32> {
+func @fold_reshape_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_collapse_of_collapse_trailing_unit_dims_dynamic
+//       CHECK: func @fold_reshape_trailing_unit_dims_dynamic
 //       CHECK: tensor.collapse_shape %{{.*}} {{\[}}[0, 1, 2, 3, 4, 5]]
 //  CHECK-SAME:   tensor<1x1x?x1x1x1xf32> into tensor<?xf32>
 
 // -----
 
-func @fold_collapse_of_expand_trailing_unit_dims(%arg0: tensor<12x42x1x1xf32>)
+func @fold_reshape_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_collapse_of_expand_trailing_unit_dims
+//       CHECK: func @fold_reshape_trailing_unit_dims
 //       CHECK: tensor.collapse_shape %{{.*}} {{\[}}[0], [1, 2, 3]]
 //  CHECK-SAME:   tensor<12x42x1x1xf32> into tensor<12x42xf32>
 
 // -----
 
-func @fold_collapse_of_expand_unit_dims_in_middle(%arg0 : tensor<?x?x?xf32>)
-    -> tensor<?x?xf32> {
+func @fold_reshapes_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_collapse_of_expand_unit_dims_in_middle
+// CHECK-LABEL: func @fold_reshapes_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_collapse_of_expand_incompatible(%arg0 : tensor<4x6x8xf32>)
+func @no_fold_reshape_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_collapse_of_expand_incompatible
+// CHECK-LABEL: func @no_fold_reshape_incompatible
 //       CHECK:   tensor.expand_shape
 //       CHECK:   tensor.collapse_shape
 
 // -----
 
-func @no_fold_collapse_of_expand_empty_expr(%arg0: tensor<3x2x2xf32>)
-    -> tensor<12x1xf32> {
+func @no_fold_reshape_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_collapse_of_expand_empty_expr
+//      CHECK: func @no_fold_reshape_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_same_static_shape(
+// CHECK-LABEL: func @pad_tensor_same_static_shape(
 //  CHECK-SAME:   %[[ARG0:.*]]: tensor<5x6xf32>
 //   CHECK-NOT:   tensor.pad
 //       CHECK:   return %[[ARG0]]
-func @pad_same_static_shape(%arg0: tensor<5x6xf32>, %a: index)
+func @pad_tensor_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_nofold_same_static_shape(
+// CHECK-LABEL: func @pad_tensor_nofold_same_static_shape(
 //  CHECK-SAME:   %[[ARG0:.*]]: tensor<5x6xf32>
 //       CHECK:   %[[PAD:.*]] = tensor.pad
 //       CHECK:   return %[[PAD]]
-func @pad_nofold_same_static_shape(%arg0: tensor<5x6xf32>, %a: index)
+func @pad_tensor_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_after_cast_different_shape(
+// CHECK-LABEL:   func @pad_tensor_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_after_cast_different_shape(%arg0: tensor<?x64x?x?xf32>)
+func @pad_tensor_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_after_cast_same_shape(
+// CHECK-LABEL:   func @pad_tensor_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_after_cast_same_shape(%arg0: tensor<?x64x?x?xf32>, %padding : index)
+func @pad_tensor_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_of_cast(
+// CHECK-LABEL: func @pad_tensor_of_cast(
 // CHECK-NOT:     tensor.cast
 // CHECK:         tensor.pad
 // CHECK:         tensor<8x?xf32> to tensor<8x32xf32>
-func @pad_of_cast(%t: tensor<8x?xf32>, %s: index) -> tensor<8x32xf32> {
+func @pad_tensor_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 @pad_cast_fold(%arg0: tensor<4x4xf32>) -> tensor<4x4xf32> {
+func @tensor_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: @pad_cast
+// CHECK-LABEL: @tensor_pad_cast
 // CHECK-SAME: %[[ARG0:.+]]: tensor<4x4xf32>
 // CHECK: return %[[ARG0]]
 
 // -----
 
-// CHECK-LABEL: func @fold_pad_source_cast(
+// CHECK-LABEL: func @fold_pad_tensor_source_cast(
 //  CHECK-SAME:                  %[[ARG0:.*]]: tensor<4x?xf32>
 //   CHECK-NOT:   tensor.cast
 //       CHECK:   %[[RESULT:.*]] = tensor.pad %[[ARG0]]
-func @fold_pad_source_cast(%arg0: tensor<4x?xf32>) -> tensor<4x4xf32> {
+func @fold_pad_tensor_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>
 }