[Mlir-commits] [mlir] [mlir][tensor] Add more tensor.extract_slice canonicalization (PR #212974)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu Jul 30 02:41:17 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Tuomas Kärnä (tkarna)
<details>
<summary>Changes</summary>
Adds more canonicalization patterns for tensor.extract_slice op.
### 1. Fold full-slice rank-reducing extract of expand_shape
Before:
%expanded = tensor.expand_shape %src [[0, 1]] output_shape [4096, 1]
: tensor<4096xf32> into tensor<4096x1xf32>
%slice = tensor.extract_slice %expanded[0, 0] [4096, 1] [1, 1]
: tensor<4096x1xf32> to tensor<4096xf32>
After:
// %slice replaced by:
%src : tensor<4096xf32>
### 2. Fold extract_slice of tensor.empty to a smaller tensor.empty
Before:
%empty = tensor.empty() : tensor<4096x1xf32>
%slice = tensor.extract_slice %empty[0, 0] [4096, 1] [1, 1]
: tensor<4096x1xf32> to tensor<4096xf32>
After:
%new_empty = tensor.empty() : tensor<4096xf32>
### 3. Fold extract_slice of linalg.fill over tensor.empty by shrinking destination
Before:
%empty = tensor.empty() : tensor<64x96xf32>
%filled = linalg.fill ins(%cst : f32) outs(%empty : tensor<64x96xf32>) -> tensor<64x96xf32>
%slice = tensor.extract_slice %filled[0, 0] [32, 48] [1, 1]
: tensor<64x96xf32> to tensor<32x48xf32>
After:
%new_empty = tensor.empty() : tensor<32x48xf32>
%new_filled = linalg.fill ins(%cst : f32) outs(%new_empty : tensor<32x48xf32>) -> tensor<32x48xf32>
Assisted-by: GPT-5.3-Codex
---
Full diff: https://github.com/llvm/llvm-project/pull/212974.diff
4 Files Affected:
- (modified) mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp (+36-1)
- (modified) mlir/lib/Dialect/Tensor/IR/TensorOps.cpp (+68-1)
- (modified) mlir/test/Dialect/Linalg/canonicalize.mlir (+30)
- (modified) mlir/test/Dialect/Tensor/canonicalize.mlir (+43)
``````````diff
diff --git a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
index 170e1edf8a55d..7240d38f24c6d 100644
--- a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
@@ -973,6 +973,41 @@ struct FoldFillWithTensorExtract : public OpRewritePattern<tensor::ExtractOp> {
}
};
+/// Fold tensor.extract_slice(linalg.fill(..., tensor.empty)) by shrinking the
+/// tensor.empty and rebuilding linalg.fill on top of it.
+struct FoldExtractSliceOfFillOfEmpty
+ : public OpRewritePattern<tensor::ExtractSliceOp> {
+public:
+ using OpRewritePattern<tensor::ExtractSliceOp>::OpRewritePattern;
+
+ LogicalResult matchAndRewrite(tensor::ExtractSliceOp extractSliceOp,
+ PatternRewriter &rewriter) const override {
+ // See if tensor input of tensor.extract_slice op is the result of a
+ // linalg.fill op.
+ auto fillOp = extractSliceOp.getSource().getDefiningOp<FillOp>();
+ if (!fillOp)
+ return failure();
+
+ // Ensure the fill op has a single use.
+ if (!fillOp->hasOneUse())
+ return failure();
+
+ // See if the output of the fill op is created by a tensor.empty op.
+ if (!fillOp.getOutputs()[0].getDefiningOp<tensor::EmptyOp>())
+ return failure();
+
+ // Create a new tensor.empty op with the smaller size of the extract_slice.
+ Value smallerEmpty = tensor::EmptyOp::create(
+ rewriter, extractSliceOp.getLoc(), extractSliceOp.getType(),
+ extractSliceOp.getSizes());
+ // Create a new linalg.fill op with the same value and the smaller empty.
+ auto newFill = FillOp::create(rewriter, extractSliceOp.getLoc(),
+ fillOp.getInputs(), smallerEmpty);
+ rewriter.replaceOp(extractSliceOp, newFill.getResult(0));
+ return success();
+ }
+};
+
/// Folds pack(fill) into a single fill op if
/// 1. The pack op does not have padding value, or
/// 2. The filled value and padding value are the same.
@@ -1103,7 +1138,7 @@ struct FoldConcatsOfFill : public OpRewritePattern<tensor::ConcatOp> {
void FillOp::getCanonicalizationPatterns(RewritePatternSet &results,
MLIRContext *context) {
results.add<FoldConcatsOfFill, FoldFillWithCopy, FoldFillWithTensorExtract,
- FoldFillWithPack, FoldFillWithPad,
+ FoldExtractSliceOfFillOfEmpty, FoldFillWithPack, FoldFillWithPad,
FoldFillWithTensorReshape<tensor::CollapseShapeOp>,
FoldFillWithTensorReshape<tensor::ExpandShapeOp>,
FoldInsertPadIntoFill, FoldFillWithTranspose>(context);
diff --git a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
index 637366a289ac9..998f8595a5f7c 100644
--- a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
+++ b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
@@ -2612,6 +2612,72 @@ class ExtractSliceOpCastFolder final : public OpRewritePattern<ExtractSliceOp> {
}
};
+/// Fold a full-slice rank-reducing extract_slice of an expand_shape back to
+/// the expand_shape source when the expanded and sliced dimensions match.
+///
+/// Example:
+/// ```
+/// %expanded = tensor.expand_shape %src [[0, 1]] output_shape [4096, 1]
+/// : tensor<4096xf32> into tensor<4096x1xf32>
+/// %slice = tensor.extract_slice %expanded[0, 0] [4096, 1] [1, 1]
+/// : tensor<4096x1xf32> to tensor<4096xf32>
+/// ```
+///
+class FoldExtractSliceOfExpandShape final
+ : public OpRewritePattern<ExtractSliceOp> {
+public:
+ using OpRewritePattern<ExtractSliceOp>::OpRewritePattern;
+
+ LogicalResult matchAndRewrite(ExtractSliceOp sliceOp,
+ PatternRewriter &rewriter) const override {
+ auto expandOp = sliceOp.getSource().getDefiningOp<ExpandShapeOp>();
+ if (!expandOp)
+ return failure();
+
+ if (sliceOp.getType() != expandOp.getSrcType())
+ return failure();
+
+ SmallVector<OpFoldResult> mixedExpandedSizes = expandOp.getMixedOutputShape();
+ if (mixedExpandedSizes.size() != sliceOp.getMixedSizes().size())
+ return failure();
+
+ for (auto [offset, size, stride, expandedSize] : llvm::zip_equal(
+ sliceOp.getMixedOffsets(), sliceOp.getMixedSizes(),
+ sliceOp.getMixedStrides(), mixedExpandedSizes)) {
+ if (getConstantIntValue(offset) != static_cast<int64_t>(0) ||
+ getConstantIntValue(stride) != static_cast<int64_t>(1))
+ return failure();
+ if (size != expandedSize)
+ return failure();
+ }
+
+ rewriter.replaceOp(sliceOp, expandOp.getSrc());
+ return success();
+ }
+};
+
+/// Fold extract_slice of tensor.empty to a smaller tensor.empty.
+class FoldExtractSliceOfEmpty final
+ : public OpRewritePattern<ExtractSliceOp> {
+public:
+ using OpRewritePattern<ExtractSliceOp>::OpRewritePattern;
+
+ LogicalResult matchAndRewrite(ExtractSliceOp sliceOp,
+ PatternRewriter &rewriter) const override {
+ auto makeSmallerEmpty = [&]() -> Value {
+ return EmptyOp::create(rewriter, sliceOp.getLoc(), sliceOp.getType(),
+ sliceOp.getSizes())
+ .getResult();
+ };
+
+ if (sliceOp.getSource().getDefiningOp<EmptyOp>()) {
+ rewriter.replaceOp(sliceOp, makeSmallerEmpty());
+ return success();
+ }
+ return failure();
+ }
+};
+
/// Slice elements from `values` into `outValues`. `counts` represents the
/// numbers of elements to stride in the original values for each dimension.
/// The output values can be used to construct a DenseElementsAttr.
@@ -2767,7 +2833,8 @@ void ExtractSliceOp::getCanonicalizationPatterns(RewritePatternSet &results,
MLIRContext *context) {
results.add<
OpWithOffsetSizesAndStridesConstantArgumentFolder<
- ExtractSliceOp, SliceReturnTypeCanonicalizer, SliceCanonicalizer>,
+ ExtractSliceOp, SliceReturnTypeCanonicalizer, SliceCanonicalizer>,
+ FoldExtractSliceOfEmpty, FoldExtractSliceOfExpandShape,
ExtractSliceOpCastFolder>(context);
}
diff --git a/mlir/test/Dialect/Linalg/canonicalize.mlir b/mlir/test/Dialect/Linalg/canonicalize.mlir
index bb11ce0d4dfb8..074f662d34245 100644
--- a/mlir/test/Dialect/Linalg/canonicalize.mlir
+++ b/mlir/test/Dialect/Linalg/canonicalize.mlir
@@ -117,6 +117,36 @@ func.func @linalg_effects(
// -----
+// CHECK-LABEL: func @fold_extract_slice_of_fill_of_empty
+// CHECK-NOT: tensor.extract_slice
+// CHECK: %[[EMPTY:.*]] = tensor.empty() : tensor<4096xf32>
+// CHECK: %[[FILL:.*]] = linalg.fill ins(%[[CST:.*]] : f32) outs(%[[EMPTY]] : tensor<4096xf32>) -> tensor<4096xf32>
+// CHECK: return %[[FILL]] : tensor<4096xf32>
+func.func @fold_extract_slice_of_fill_of_empty(%cst : f32) -> tensor<4096xf32> {
+ %empty = tensor.empty() : tensor<4096x1xf32>
+ %filled = linalg.fill ins(%cst : f32) outs(%empty : tensor<4096x1xf32>) -> tensor<4096x1xf32>
+ %slice = tensor.extract_slice %filled[0, 0] [4096, 1] [1, 1]
+ : tensor<4096x1xf32> to tensor<4096xf32>
+ return %slice : tensor<4096xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @fold_extract_slice_of_fill_of_empty_2d
+// CHECK-NOT: tensor.extract_slice
+// CHECK: %[[EMPTY:.*]] = tensor.empty() : tensor<32x48xf32>
+// CHECK: %[[FILL:.*]] = linalg.fill ins(%[[CST:.*]] : f32) outs(%[[EMPTY]] : tensor<32x48xf32>) -> tensor<32x48xf32>
+// CHECK: return %[[FILL]] : tensor<32x48xf32>
+func.func @fold_extract_slice_of_fill_of_empty_2d(%cst : f32) -> tensor<32x48xf32> {
+ %empty = tensor.empty() : tensor<64x96xf32>
+ %filled = linalg.fill ins(%cst : f32) outs(%empty : tensor<64x96xf32>) -> tensor<64x96xf32>
+ %slice = tensor.extract_slice %filled[0, 0] [32, 48] [1, 1]
+ : tensor<64x96xf32> to tensor<32x48xf32>
+ return %slice : tensor<32x48xf32>
+}
+
+// -----
+
#map = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
func.func @remove_no_op(%arg0 : tensor<?x?x?xf32>, %arg1 : tensor<?x?x?xf32>)
-> (tensor<?x?x?xf32>, tensor<?x?x?xf32>) {
diff --git a/mlir/test/Dialect/Tensor/canonicalize.mlir b/mlir/test/Dialect/Tensor/canonicalize.mlir
index 67b7ab99c5d18..adcfe5b97b558 100644
--- a/mlir/test/Dialect/Tensor/canonicalize.mlir
+++ b/mlir/test/Dialect/Tensor/canonicalize.mlir
@@ -623,6 +623,49 @@ func.func @trivial_slice(%arg0 : tensor<4x6x16x32xi8>) -> tensor<4x6x16x32xi8> {
// -----
+// CHECK-LABEL: func @fold_extract_slice_of_expand_shape
+// CHECK-SAME: %[[ARG0:.*]]: tensor<4096xf32>
+// CHECK-NOT: tensor.expand_shape
+// CHECK-NOT: tensor.extract_slice
+// CHECK: return %[[ARG0]] : tensor<4096xf32>
+func.func @fold_extract_slice_of_expand_shape(
+ %arg0 : tensor<4096xf32>) -> tensor<4096xf32> {
+ %expanded = tensor.expand_shape %arg0 [[0, 1]] output_shape [4096, 1]
+ : tensor<4096xf32> into tensor<4096x1xf32>
+ %slice = tensor.extract_slice %expanded[0, 0] [4096, 1] [1, 1]
+ : tensor<4096x1xf32> to tensor<4096xf32>
+ return %slice : tensor<4096xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @dont_fold_extract_slice_of_expand_shape_with_different_sizes
+// CHECK: tensor.expand_shape
+// CHECK: tensor.extract_slice
+func.func @dont_fold_extract_slice_of_expand_shape_with_different_sizes(
+ %arg0 : tensor<4096xf32>) -> tensor<1024xf32> {
+ %expanded = tensor.expand_shape %arg0 [[0, 1]] output_shape [4096, 1]
+ : tensor<4096xf32> into tensor<4096x1xf32>
+ %slice = tensor.extract_slice %expanded[0, 0] [1024, 1] [1, 1]
+ : tensor<4096x1xf32> to tensor<1024xf32>
+ return %slice : tensor<1024xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @fold_extract_slice_of_empty
+// CHECK-NOT: tensor.extract_slice
+// CHECK: %[[EMPTY:.*]] = tensor.empty() : tensor<4096xf32>
+// CHECK: return %[[EMPTY]] : tensor<4096xf32>
+func.func @fold_extract_slice_of_empty() -> tensor<4096xf32> {
+ %empty = tensor.empty() : tensor<4096x1xf32>
+ %slice = tensor.extract_slice %empty[0, 0] [4096, 1] [1, 1]
+ : tensor<4096x1xf32> to tensor<4096xf32>
+ return %slice : tensor<4096xf32>
+}
+
+// -----
+
// CHECK-LABEL: func @trivial_insert_slice
// CHECK-SAME: %[[ARG0:.[a-z0-9A-Z_]+]]: tensor<4x6x16x32xi8>
// CHECK-NOT: tensor.extract_slice
``````````
</details>
https://github.com/llvm/llvm-project/pull/212974
More information about the Mlir-commits
mailing list