[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

220

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

275

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

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

1711

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

1752

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

[nicolasvasilache]

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

242

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

371

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

242

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

242

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

242

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

242

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

return reshapeLikeShapesAreCompatible(
       [&](const Twine &msg) { return op->emitOpError(msg); },
       collapsedType.getShape(), expandedType.getShape(),
       op.getReassociationIndices(), isExpandingReshape);
 }
 
 /// 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();
     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 intermediateType = collapseOp.getSrcType();
+    ShapedType resultType = collapseOp.getResultType();
+
+    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();
+    }
+
+    int 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<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 intermediateType = expandOp.getSrcType();
+    ShapedType resultType = expandOp.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);
+        getReassociationIndicesForReshape(srcType, resultType);
     if (!reassociationIndices)
       return failure();
-    bool originalOpExpands =
-        intermediateType.getRank() > srcReshapeSrcType.getRank();
-    bool resultingOpExpands =
-        resultType.getRank() > srcReshapeSrcType.getRank();
+    bool originalOpExpands = intermediateType.getRank() > srcType.getRank();
+    bool resultingOpExpands = resultType.getRank() > srcType.getRank();
     if (!(resultingOpExpands ^ originalOpExpands))
-      rewriter.replaceOpWithNewOp<InverseReshapeOpTy>(
-          reshapeOp, resultType, srcReshapeOp.src(), *reassociationIndices);
+      rewriter.replaceOpWithNewOp<CollapseOpTy>(
+          expandOp, resultType, collapseOp.src(), *reassociationIndices);
     else
-      rewriter.replaceOpWithNewOp<ReshapeOpTy>(
-          reshapeOp, resultType, srcReshapeOp.src(), *reassociationIndices);
+      rewriter.replaceOpWithNewOp<ExpandOpTy>(
+          expandOp, resultType, collapseOp.src(), *reassociationIndices);
     return success();
   }
 };
 
 } // namespace mlir
 
 #endif // MLIR_DIALECT_UTILS_RESHAPEOPSUTILS_H

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

 }
 
 LogicalResult ExpandShapeOp::verify() {
   return verifyReshapeOp(*this, getResultType(), getSrcType());
 }
 
 void ExpandShapeOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                                 MLIRContext *context) {
-  results.add<CollapseReshapeOps<ExpandShapeOp>,
-              CollapseMixedReshapeOps<ExpandShapeOp, CollapseShapeOp>>(context);
+  results.add<ComposeReassociativeReshapeOps<ExpandShapeOp>,
+              ComposeExpandOfCollapseOp<ExpandShapeOp, CollapseShapeOp>>(
+      context);
 }
 
 LogicalResult CollapseShapeOp::verify() {
   return verifyReshapeOp(*this, getSrcType(), getResultType());
 }
 
 struct CollapseShapeOpMemRefCastFolder
     : public OpRewritePattern<CollapseShapeOp> {
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) {
   return foldReshapeOp<CollapseShapeOp, ExpandShapeOp>(*this, 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/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 @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>
 }