[Mlir-commits] [mlir] [mlir][tensor] Add more tensor.extract_slice canonicalization (PR #212974)
Tuomas Kärnä
llvmlistbot at llvm.org
Thu Jul 30 03:27:49 PDT 2026
https://github.com/tkarna updated https://github.com/llvm/llvm-project/pull/212974
>From 172562c60019474a4505c51b14123a9e9174bd44 Mon Sep 17 00:00:00 2001
From: Tuomas Karna <tuomas.karna at intel.com>
Date: Wed, 29 Jul 2026 12:09:53 +0300
Subject: [PATCH 1/4] tensor: add FoldExtractSliceOfExpandShape
canonicalization pattern
---
mlir/lib/Dialect/Tensor/IR/TensorOps.cpp | 48 +++++++++++++++++++++-
mlir/test/Dialect/Tensor/canonicalize.mlir | 30 ++++++++++++++
2 files changed, 76 insertions(+), 2 deletions(-)
diff --git a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
index 637366a289ac9..40d2780306b1f 100644
--- a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
+++ b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
@@ -2612,6 +2612,50 @@ 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();
+ }
+};
+
/// 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,8 +2811,8 @@ void ExtractSliceOp::getCanonicalizationPatterns(RewritePatternSet &results,
MLIRContext *context) {
results.add<
OpWithOffsetSizesAndStridesConstantArgumentFolder<
- ExtractSliceOp, SliceReturnTypeCanonicalizer, SliceCanonicalizer>,
- ExtractSliceOpCastFolder>(context);
+ ExtractSliceOp, SliceReturnTypeCanonicalizer, SliceCanonicalizer>,
+ FoldExtractSliceOfExpandShape, ExtractSliceOpCastFolder>(context);
}
//
diff --git a/mlir/test/Dialect/Tensor/canonicalize.mlir b/mlir/test/Dialect/Tensor/canonicalize.mlir
index 67b7ab99c5d18..dfff209187cf1 100644
--- a/mlir/test/Dialect/Tensor/canonicalize.mlir
+++ b/mlir/test/Dialect/Tensor/canonicalize.mlir
@@ -623,6 +623,36 @@ 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 @trivial_insert_slice
// CHECK-SAME: %[[ARG0:.[a-z0-9A-Z_]+]]: tensor<4x6x16x32xi8>
// CHECK-NOT: tensor.extract_slice
>From 7d89680cb36b5908cdf9c623c6ef2496164cfc07 Mon Sep 17 00:00:00 2001
From: Tuomas Karna <tuomas.karna at intel.com>
Date: Wed, 29 Jul 2026 12:50:23 +0300
Subject: [PATCH 2/4] tensor: FoldExtractSliceOfEmpty canonicalization pattern
---
mlir/lib/Dialect/Tensor/IR/TensorOps.cpp | 25 +++++++++++++++++++++-
mlir/test/Dialect/Tensor/canonicalize.mlir | 13 +++++++++++
2 files changed, 37 insertions(+), 1 deletion(-)
diff --git a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
index 40d2780306b1f..998f8595a5f7c 100644
--- a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
+++ b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
@@ -2656,6 +2656,28 @@ class FoldExtractSliceOfExpandShape final
}
};
+/// 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.
@@ -2812,7 +2834,8 @@ void ExtractSliceOp::getCanonicalizationPatterns(RewritePatternSet &results,
results.add<
OpWithOffsetSizesAndStridesConstantArgumentFolder<
ExtractSliceOp, SliceReturnTypeCanonicalizer, SliceCanonicalizer>,
- FoldExtractSliceOfExpandShape, ExtractSliceOpCastFolder>(context);
+ FoldExtractSliceOfEmpty, FoldExtractSliceOfExpandShape,
+ ExtractSliceOpCastFolder>(context);
}
//
diff --git a/mlir/test/Dialect/Tensor/canonicalize.mlir b/mlir/test/Dialect/Tensor/canonicalize.mlir
index dfff209187cf1..adcfe5b97b558 100644
--- a/mlir/test/Dialect/Tensor/canonicalize.mlir
+++ b/mlir/test/Dialect/Tensor/canonicalize.mlir
@@ -653,6 +653,19 @@ func.func @dont_fold_extract_slice_of_expand_shape_with_different_sizes(
// -----
+// 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
>From 8fdf5f5a79da709f9e467d5e2867c1382cd32b96 Mon Sep 17 00:00:00 2001
From: Tuomas Karna <tuomas.karna at intel.com>
Date: Wed, 29 Jul 2026 14:56:41 +0300
Subject: [PATCH 3/4] linalg: add FoldExtractSliceOfFillOfEmpty
canonicalization pattern
---
mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp | 37 +++++++++++++++++++++-
mlir/test/Dialect/Linalg/canonicalize.mlir | 30 ++++++++++++++++++
2 files changed, 66 insertions(+), 1 deletion(-)
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/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>) {
>From 156d359a4608864feffc258eb4d32c64e074fdb9 Mon Sep 17 00:00:00 2001
From: Tuomas Karna <tuomas.karna at intel.com>
Date: Thu, 30 Jul 2026 13:04:13 +0300
Subject: [PATCH 4/4] FoldExtractSliceOfEmpty only applies to rank-reducing
extract slice
---
mlir/lib/Dialect/Tensor/IR/TensorOps.cpp | 58 +++++++++++++++++-------
1 file changed, 41 insertions(+), 17 deletions(-)
diff --git a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
index 998f8595a5f7c..b357cb8dce50e 100644
--- a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
+++ b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
@@ -2637,13 +2637,14 @@ class FoldExtractSliceOfExpandShape final
if (sliceOp.getType() != expandOp.getSrcType())
return failure();
- SmallVector<OpFoldResult> mixedExpandedSizes = expandOp.getMixedOutputShape();
+ 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)) {
+ 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();
@@ -2656,25 +2657,48 @@ class FoldExtractSliceOfExpandShape final
}
};
-/// Fold extract_slice of tensor.empty to a smaller tensor.empty.
-class FoldExtractSliceOfEmpty final
- : public OpRewritePattern<ExtractSliceOp> {
+/// Fold a rank-reducing no-op extract_slice of a tensor.empty into a smaller
+/// tensor.empty.
+///
+/// The slice must be a full/identity slice: all offsets are 0, all strides are
+/// 1, and each size matches the corresponding tensor.empty source dimension.
+/// The slice must also be rank-reducing. Restricting to this case avoids
+/// undoing transforms that intentionally allocate a larger tensor.empty and
+/// then slice it (e.g. transform.tensor.make_loop_independent).
+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>())
+ return failure();
- if (sliceOp.getSource().getDefiningOp<EmptyOp>()) {
- rewriter.replaceOp(sliceOp, makeSmallerEmpty());
- return success();
+ // Only fold rank-reducing slices. Non-rank-reducing identity slices are
+ // already handled by ExtractSliceOp::fold.
+ if (sliceOp.getType().getRank() >= sliceOp.getSourceType().getRank())
+ return failure();
+
+ // Only fold full/identity slices: all offsets are 0, all strides are 1,
+ // and each size matches the corresponding (static) source dimension.
+ ArrayRef<int64_t> sourceShape = sliceOp.getSourceType().getShape();
+ for (auto [offset, size, stride, srcDim] :
+ llvm::zip_equal(sliceOp.getMixedOffsets(), sliceOp.getMixedSizes(),
+ sliceOp.getMixedStrides(), sourceShape)) {
+ if (getConstantIntValue(offset) != static_cast<int64_t>(0) ||
+ getConstantIntValue(stride) != static_cast<int64_t>(1))
+ return failure();
+ // Bail out if the source dim is dynamic or the size does not provably
+ // match it.
+ if (ShapedType::isDynamic(srcDim) || getConstantIntValue(size) != srcDim)
+ return failure();
}
- return failure();
+
+ rewriter.replaceOp(sliceOp,
+ EmptyOp::create(rewriter, sliceOp.getLoc(),
+ sliceOp.getType(), sliceOp.getSizes())
+ .getResult());
+ return success();
}
};
@@ -2833,7 +2857,7 @@ void ExtractSliceOp::getCanonicalizationPatterns(RewritePatternSet &results,
MLIRContext *context) {
results.add<
OpWithOffsetSizesAndStridesConstantArgumentFolder<
- ExtractSliceOp, SliceReturnTypeCanonicalizer, SliceCanonicalizer>,
+ ExtractSliceOp, SliceReturnTypeCanonicalizer, SliceCanonicalizer>,
FoldExtractSliceOfEmpty, FoldExtractSliceOfExpandShape,
ExtractSliceOpCastFolder>(context);
}
More information about the Mlir-commits
mailing list