[Mlir-commits] [mlir] [mlir][tensor] Consolidate tensor fold patterns and rename related file (PR #192820)
Longsheng Mou
llvmlistbot at llvm.org
Sat Apr 18 20:09:39 PDT 2026
https://github.com/CoTinker created https://github.com/llvm/llvm-project/pull/192820
This PR moves `MergeConsecutiveExtractSlice` from `MergeConsecutiveInsertExtractSlicePatterns.cpp` to `FoldTensorSubsetOps.cpp`, and removes the duplicate `MergeConsecutiveInsertSlice` pattern in favor of `InsertSliceOfInsertSliceFolder`, which already exists in `FoldTensorSubsetOps.cpp` and provides equivalent functionality with greater stability. Since the merge-related patterns have been fully migrated out, `MergeConsecutiveInsertExtractSlicePatterns.cpp` is renamed to `DropRedundantRankExpansionPatterns.cpp` to better reflect its remaining responsibilities.
>From 5de8634b7f839b082fd16b527748e1206b850b7d Mon Sep 17 00:00:00 2001
From: Longsheng Mou <longshengmou at gmail.com>
Date: Sat, 18 Apr 2026 18:28:23 +0800
Subject: [PATCH] [mlir][tensor] Consolidate fold patterns in
FoldTensorSubsetOps.cpp
This PR moves `MergeConsecutiveExtractSlice` from `MergeConsecutiveInsertExtractSlicePatterns.cpp` to `FoldTensorSubsetOps.cpp`, and removes the duplicate `MergeConsecutiveInsertSlice` pattern in favor of `InsertSliceOfInsertSliceFolder`, which already exists in `FoldTensorSubsetOps.cpp` and provides equivalent functionality with greater stability. Since the merge-related patterns have been fully migrated out, `MergeConsecutiveInsertExtractSlicePatterns.cpp` is renamed to `DropRedundantRankExpansionPatterns.cpp` to better reflect its remaining responsibilities.
---
.../Dialect/Tensor/Transforms/CMakeLists.txt | 2 +-
...=> DropRedundantRankExpansionPatterns.cpp} | 70 +------------------
.../Tensor/Transforms/FoldTensorSubsetOps.cpp | 41 +++++++++--
...fold-consecutive-insert-extract-slice.mlir | 2 +-
.../Tensor/fold-tensor-subset-ops.mlir | 58 +++++++++++++++
5 files changed, 96 insertions(+), 77 deletions(-)
rename mlir/lib/Dialect/Tensor/Transforms/{MergeConsecutiveInsertExtractSlicePatterns.cpp => DropRedundantRankExpansionPatterns.cpp} (71%)
diff --git a/mlir/lib/Dialect/Tensor/Transforms/CMakeLists.txt b/mlir/lib/Dialect/Tensor/Transforms/CMakeLists.txt
index 99e1c4fec8467..33d32c592a844 100644
--- a/mlir/lib/Dialect/Tensor/Transforms/CMakeLists.txt
+++ b/mlir/lib/Dialect/Tensor/Transforms/CMakeLists.txt
@@ -1,11 +1,11 @@
add_mlir_dialect_library(MLIRTensorTransforms
BufferizableOpInterfaceImpl.cpp
ConcatOpPatterns.cpp
+ DropRedundantRankExpansionPatterns.cpp
EmptyOpPatterns.cpp
ExtractSliceFromReshapeUtils.cpp
FoldTensorSubsetOps.cpp
IndependenceTransforms.cpp
- MergeConsecutiveInsertExtractSlicePatterns.cpp
ReshapePatterns.cpp
RewriteAsConstant.cpp
RuntimeOpVerification.cpp
diff --git a/mlir/lib/Dialect/Tensor/Transforms/MergeConsecutiveInsertExtractSlicePatterns.cpp b/mlir/lib/Dialect/Tensor/Transforms/DropRedundantRankExpansionPatterns.cpp
similarity index 71%
rename from mlir/lib/Dialect/Tensor/Transforms/MergeConsecutiveInsertExtractSlicePatterns.cpp
rename to mlir/lib/Dialect/Tensor/Transforms/DropRedundantRankExpansionPatterns.cpp
index ff003e486d21c..4253548d11f49 100644
--- a/mlir/lib/Dialect/Tensor/Transforms/MergeConsecutiveInsertExtractSlicePatterns.cpp
+++ b/mlir/lib/Dialect/Tensor/Transforms/DropRedundantRankExpansionPatterns.cpp
@@ -1,4 +1,4 @@
-//===- MergeConsecutiveInsertExtractSlicePatterns.cpp ---------------------===//
+//===- DropRedundantRankExpansionPatterns.cpp -----------------------------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
@@ -18,66 +18,6 @@ using namespace mlir;
using namespace mlir::tensor;
namespace {
-/// Merges consecutive tensor.extract_slice ops into one.
-// TODO: move to FoldTensorSubsetOps and unify APIs with FoldMemRefAliasOps.
-struct MergeConsecutiveExtractSlice : public OpRewritePattern<ExtractSliceOp> {
- using OpRewritePattern::OpRewritePattern;
-
- LogicalResult matchAndRewrite(ExtractSliceOp nextOp,
- PatternRewriter &rewriter) const override {
- auto prevOp = nextOp.getSource().getDefiningOp<ExtractSliceOp>();
- if (!prevOp)
- return failure();
-
- SmallVector<OpFoldResult> newOffsets, newSizes, newStrides;
- if (failed(affine::mergeOffsetsSizesAndStrides(
- rewriter, nextOp.getLoc(), prevOp, nextOp, prevOp.getDroppedDims(),
- newOffsets, newSizes, newStrides)))
- return failure();
-
- rewriter.replaceOpWithNewOp<ExtractSliceOp>(nextOp, nextOp.getType(),
- prevOp.getSource(), newOffsets,
- newSizes, newStrides);
- return success();
- }
-};
-
-/// Merges consecutive tensor.insert_slice ops into one.
-// TODO: move to FoldTensorSubsetOps and unify APIs with FoldMemRefAliasOps.
-template <typename OpTy>
-struct MergeConsecutiveInsertSlice : public OpRewritePattern<OpTy> {
- using OpRewritePattern<OpTy>::OpRewritePattern;
-
- LogicalResult matchAndRewrite(OpTy nextOp,
- PatternRewriter &rewriter) const override {
- auto prevOp = nextOp.getSource().template getDefiningOp<InsertSliceOp>();
- if (!prevOp)
- return failure();
-
- if (!prevOp.hasUnitStride() || !nextOp.hasUnitStride())
- return failure();
-
- // The first insert_slice op should be rank reducing to make sure we cover
- // the full source tensor to be inserted in the second insert_slice op.
- SliceVerificationResult result =
- isRankReducedType(prevOp.getDestType(), prevOp.getSourceType());
- if (result != SliceVerificationResult::Success)
- return failure();
-
- // Dynamic dimensions can pass rank reducing check in the above, e.g,
- // inserting <?xf32> into <1x?x1xf32>. For such cases we cannot be certain
- // the dynamic size covers the full tensor.
- if (!prevOp.getSourceType().hasStaticShape() ||
- !prevOp.getDestType().hasStaticShape())
- return failure();
-
- rewriter.replaceOpWithNewOp<OpTy>(
- nextOp, prevOp.getSource(), nextOp.getDest(), nextOp.getMixedOffsets(),
- nextOp.getMixedSizes(), nextOp.getMixedStrides());
- return success();
- }
-};
-
/// Drop redundant rank expansion of insert_slice that are directly followed
/// by extract_slice. E.g.:
/// %0 = tensor.insert_slice ... : tensor<5x10xf32> into tensor<1x1x5x10xf32>
@@ -227,14 +167,6 @@ struct DropRedundantRankExpansionOnInsertSliceOfExtractSlice final
};
} // namespace
-void mlir::tensor::populateMergeConsecutiveInsertExtractSlicePatterns(
- RewritePatternSet &patterns) {
- patterns.add<MergeConsecutiveExtractSlice,
- MergeConsecutiveInsertSlice<InsertSliceOp>,
- MergeConsecutiveInsertSlice<ParallelInsertSliceOp>>(
- patterns.getContext());
-}
-
void mlir::tensor::populateDropRedundantInsertSliceRankExpansionPatterns(
RewritePatternSet &patterns) {
patterns.add<DropRedundantRankExpansionOnExtractSliceOfInsertSlice,
diff --git a/mlir/lib/Dialect/Tensor/Transforms/FoldTensorSubsetOps.cpp b/mlir/lib/Dialect/Tensor/Transforms/FoldTensorSubsetOps.cpp
index 14f96be5b56dd..5f838eae95de9 100644
--- a/mlir/lib/Dialect/Tensor/Transforms/FoldTensorSubsetOps.cpp
+++ b/mlir/lib/Dialect/Tensor/Transforms/FoldTensorSubsetOps.cpp
@@ -220,12 +220,28 @@ struct InsertSliceOfInsertSliceFolder : public OpRewritePattern<OpTy> {
}
};
-void tensor::populateFoldTensorSubsetOpPatterns(RewritePatternSet &patterns) {
- populateFoldTensorSubsetIntoVectorTransferPatterns(patterns);
- patterns.add<InsertSliceOfInsertSliceFolder<tensor::InsertSliceOp>,
- InsertSliceOfInsertSliceFolder<tensor::ParallelInsertSliceOp>>(
- patterns.getContext());
-}
+struct MergeConsecutiveExtractSlice
+ : public OpRewritePattern<tensor::ExtractSliceOp> {
+ using OpRewritePattern::OpRewritePattern;
+
+ LogicalResult matchAndRewrite(tensor::ExtractSliceOp nextOp,
+ PatternRewriter &rewriter) const override {
+ auto prevOp = nextOp.getSource().getDefiningOp<tensor::ExtractSliceOp>();
+ if (!prevOp)
+ return failure();
+
+ SmallVector<OpFoldResult> newOffsets, newSizes, newStrides;
+ if (failed(affine::mergeOffsetsSizesAndStrides(
+ rewriter, nextOp.getLoc(), prevOp, nextOp, prevOp.getDroppedDims(),
+ newOffsets, newSizes, newStrides)))
+ return failure();
+
+ rewriter.replaceOpWithNewOp<tensor::ExtractSliceOp>(
+ nextOp, nextOp.getType(), prevOp.getSource(), newOffsets, newSizes,
+ newStrides);
+ return success();
+ }
+};
void tensor::populateFoldTensorSubsetIntoVectorTransferPatterns(
RewritePatternSet &patterns) {
@@ -233,6 +249,19 @@ void tensor::populateFoldTensorSubsetIntoVectorTransferPatterns(
InsertSliceOfTransferWriteOpFolder>(patterns.getContext());
}
+void tensor::populateMergeConsecutiveInsertExtractSlicePatterns(
+ RewritePatternSet &patterns) {
+ patterns.add<MergeConsecutiveExtractSlice,
+ InsertSliceOfInsertSliceFolder<tensor::InsertSliceOp>,
+ InsertSliceOfInsertSliceFolder<tensor::ParallelInsertSliceOp>>(
+ patterns.getContext());
+}
+
+void tensor::populateFoldTensorSubsetOpPatterns(RewritePatternSet &patterns) {
+ populateFoldTensorSubsetIntoVectorTransferPatterns(patterns);
+ populateMergeConsecutiveInsertExtractSlicePatterns(patterns);
+}
+
//===----------------------------------------------------------------------===//
// Pass registration
//===----------------------------------------------------------------------===//
diff --git a/mlir/test/Dialect/Tensor/fold-consecutive-insert-extract-slice.mlir b/mlir/test/Dialect/Tensor/fold-consecutive-insert-extract-slice.mlir
index 750a8d0edf0e2..fb1d739862096 100644
--- a/mlir/test/Dialect/Tensor/fold-consecutive-insert-extract-slice.mlir
+++ b/mlir/test/Dialect/Tensor/fold-consecutive-insert-extract-slice.mlir
@@ -80,7 +80,7 @@ func.func @insert_slice_rank_reducing_dynamic_shape(
}
// CHECK-LABEL: func.func @insert_slice_rank_reducing_dynamic_shape
-// CHECK-COUNT-2: tensor.insert_slice
+// CHECK-COUNT-1: tensor.insert_slice
// -----
diff --git a/mlir/test/Dialect/Tensor/fold-tensor-subset-ops.mlir b/mlir/test/Dialect/Tensor/fold-tensor-subset-ops.mlir
index 45937e94f08ff..2ef7c075ad584 100644
--- a/mlir/test/Dialect/Tensor/fold-tensor-subset-ops.mlir
+++ b/mlir/test/Dialect/Tensor/fold-tensor-subset-ops.mlir
@@ -415,3 +415,61 @@ func.func @parallel_insert_slice_of_insert_slice_dynamic(
}
return %0: tensor<12x34xf32>
}
+
+// -----
+
+// CHECK: #[[$map:.*]] = affine_map<()[s0, s1] -> (s0 + s1)>
+// CHECK-LABEL: func.func @extract_slice_same_rank
+// CHECK-SAME: (%[[SOURCE:.+]]: tensor<?x?x?x?xf32>, %[[OFFSET0:.+]]: index, %[[OFFSET1:.+]]: index, %{{.+}}: index, %[[SIZE1:.+]]: index)
+// CHECK: %[[OFFSET:.+]] = affine.apply #[[$map]]()[%[[OFFSET1]], %[[OFFSET0]]]
+// CHECK: %[[EXTRACT:.+]] = tensor.extract_slice %[[SOURCE]][7, 9, 11, %[[OFFSET]]] [8, 16, 32, %[[SIZE1]]] [1, 1, 1, 1]
+// CHECK: return %[[EXTRACT]] : tensor<8x16x32x?xf32>
+func.func @extract_slice_same_rank(
+ %src: tensor<?x?x?x?xf32>, %offset0: index, %offset1: index, %size0: index, %size1: index) -> tensor<8x16x32x?xf32> {
+ %0 = tensor.extract_slice %src[0, 1, 2, %offset0] [128, 128, 128, %size0] [1, 1, 1, 1] : tensor<?x?x?x?xf32> to tensor<128x128x128x?xf32>
+ %1 = tensor.extract_slice %0[7, 8, 9, %offset1] [8, 16, 32, %size1] [1, 1, 1, 1] : tensor<128x128x128x?xf32> to tensor<8x16x32x?xf32>
+ return %1: tensor<8x16x32x?xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @extract_slice_rank_reducing_consumer
+// CHECK: tensor.extract_slice %{{.+}}[7, 9, 11, %{{.+}}] [1, 16, 1, %{{.+}}] [1, 1, 1, 1] : tensor<?x?x?x?xf32> to tensor<16x?xf32>
+func.func @extract_slice_rank_reducing_consumer(
+ %src: tensor<?x?x?x?xf32>, %offset0: index, %offset1: index, %size0: index, %size1: index) -> tensor<16x?xf32> {
+ %0 = tensor.extract_slice %src[0, 1, 2, %offset0] [128, 128, 128, %size0] [1, 1, 1, 1] : tensor<?x?x?x?xf32> to tensor<128x128x128x?xf32>
+ %1 = tensor.extract_slice %0[7, 8, 9, %offset1] [1, 16, 1, %size1] [1, 1, 1, 1] : tensor<128x128x128x?xf32> to tensor<16x?xf32>
+ return %1: tensor<16x?xf32>
+}
+
+// -----
+
+// CHECK: #[[$map:.*]] = affine_map<()[s0, s1] -> (s0 + s1)>
+// CHECK-LABEL: func.func @extract_slice_rank_reducing_producer
+// CHECK-SAME: (%[[SRC:.+]]: tensor<?x?x?x?xf32>, %[[OFFSET0:.+]]: index, %[[OFFSET1:.+]]: index, %{{.+}}: index, %[[SIZE1:.+]]: index)
+// CHECK: %[[OFFSET:.+]] = affine.apply #[[$map]]()[%[[OFFSET1]], %[[OFFSET0]]]
+// CHECK: %[[EXTRACT:.+]] = tensor.extract_slice %[[SRC]][0, 8, 2, %[[OFFSET]]] [1, 8, 1, %[[SIZE1]]] [1, 1, 1, 1] : tensor<?x?x?x?xf32> to tensor<8x?xf32>
+// CHECK: return %[[EXTRACT]] : tensor<8x?xf32>
+func.func @extract_slice_rank_reducing_producer(
+ %src: tensor<?x?x?x?xf32>, %offset0: index, %offset1: index, %size0: index, %size1: index) -> tensor<8x?xf32> {
+ %0 = tensor.extract_slice %src[0, 1, 2, %offset0] [1, 128, 1, %size0] [1, 1, 1, 1] : tensor<?x?x?x?xf32> to tensor<128x?xf32>
+ %1 = tensor.extract_slice %0[7, %offset1] [8, %size1] [1, 1] : tensor<128x?xf32> to tensor<8x?xf32>
+ return %1: tensor<8x?xf32>
+}
+
+// -----
+
+// CHECK: #[[$map_0:.+]] = affine_map<()[s0, s1, s2] -> (s0 * s1 + s2)>
+// CHECK: #[[$map_1:.+]] = affine_map<()[s0, s1] -> (s0 * s1)>
+// CHECK-LABEL: func.func @extract_slice_non_one_stride
+// CHECK-SAME: (%[[SRC:.+]]: tensor<?xf32>, %[[OFFSET0:.+]]: index, %[[OFFSET1:.+]]: index, %{{.+}}: index, %[[SIZE1:.+]]: index, %[[STRIDE0:.+]]: index, %[[STRIDE1:.+]]: index)
+// CHECK: %[[OFFSET:.+]] = affine.apply #[[$map_0]]()[%[[OFFSET1]], %[[STRIDE0]], %[[OFFSET0]]]
+// CHECK: %[[STRIDE:.+]] = affine.apply #[[$map_1]]()[%[[STRIDE1]], %[[STRIDE0]]]
+// CHECK: %[[EXTRACT:.+]] = tensor.extract_slice %[[SRC]][%[[OFFSET]]] [%[[SIZE1]]] [%[[STRIDE]]] : tensor<?xf32> to tensor<?xf32>
+// CHECK: return %[[EXTRACT]] : tensor<?xf32>
+func.func @extract_slice_non_one_stride(
+ %src: tensor<?xf32>, %offset0: index, %offset1: index, %size0: index, %size1: index, %stride0: index, %stride1: index) -> tensor<?xf32> {
+ %0 = tensor.extract_slice %src[%offset0] [%size0] [%stride0] : tensor<?xf32> to tensor<?xf32>
+ %1 = tensor.extract_slice %0[%offset1] [%size1] [%stride1] : tensor<?xf32> to tensor<?xf32>
+ return %1: tensor<?xf32>
+}
More information about the Mlir-commits
mailing list