[Mlir-commits] [mlir] [MLIR] Fix canonicalization of extract_slice(unpack) (PR #181840)
Mikhail Romanov
llvmlistbot at llvm.org
Thu Feb 19 12:05:56 PST 2026
https://github.com/Mmi257ia updated https://github.com/llvm/llvm-project/pull/181840
>From dafdb9df021d664ff7c57fd6ad17b2028de041b9 Mon Sep 17 00:00:00 2001
From: Mikhail Romanov <mmromanov at ispras.ru>
Date: Mon, 16 Feb 2026 18:33:15 +0300
Subject: [PATCH 1/2] [MLIR] Fix canonicalization of extract_slice(unpack)
This patch fixes an issue in the canonicalization pattern that folds
`extract_slice(unpack)` operation sequence. Previously, the pattern would
incorrectly apply the transformation even when the `extract_slice` was operating
on dimensions that are not the unpacked dimension, leading to invalid IR transformations.
The canonicalization pattern now checks that the `extract_slice` is only
applied to the dimension that was actually unpacked. If the slice is taken
along other dimensions, the transformation is prevented to maintain correctness.
Example of problematic case that is now handled correctly:
```
%unpack = linalg.unpack %src
inner_dims_pos = [1]
inner_tiles = [16]
into %dest : tensor<30x2x16xf32> -> tensor<30x32xf32>
%extracted_slice = tensor.extract_slice %unpack
[0, 0] [28, 32] [1, 1] : tensor<30x32xf32> to tensor<28x32xf32>
```
After canonicalization, this results in an invalid `linalg.unpack`:
```
%unpack = linalg.unpack %src
inner_dims_pos = [1]
inner_tiles = [16]
into %dest : tensor<30x2x16xf32> -> tensor<28x32xf32>
```
Signed-off-by: Mikhail Romanov <mmromanov at ispras.ru>
---
mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp | 11 ++++++
mlir/test/Dialect/Linalg/canonicalize.mlir | 44 ++++++++++++++++++++++
2 files changed, 55 insertions(+)
diff --git a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
index eba3fa6db2126..b02fe5564393a 100644
--- a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
@@ -6394,8 +6394,10 @@ bool UnPackOp::canFoldSliceOp(tensor::ExtractSliceOp sliceOp) {
RankedTensorType unpackedTypeAfterFold = sliceOp.getResultType();
SmallVector<int64_t> outerShapeWithoutTranspose =
getPackedOuterShapeWithoutTransposition(*this);
+ SmallVector<bool> areOuterDimsTiled(outerShapeWithoutTranspose.size(), false);
for (auto [pos, tileSize] :
llvm::zip_equal(this->getInnerDimsPos(), this->getStaticInnerTiles())) {
+ areOuterDimsTiled[pos] = true;
if (unpackedTypeAfterFold.isDynamicDim(pos))
return false;
if (ShapedType::isDynamic(outerShapeWithoutTranspose[pos]))
@@ -6407,6 +6409,15 @@ bool UnPackOp::canFoldSliceOp(tensor::ExtractSliceOp sliceOp) {
if (paddingSize >= tileSize)
return false;
}
+ for (int64_t pos = 0, e = outerShapeWithoutTranspose.size(); pos < e; ++pos) {
+ if (areOuterDimsTiled[pos])
+ continue;
+ int64_t dim = outerShapeWithoutTranspose[pos];
+ if (ShapedType::isDynamic(dim))
+ return false;
+ if (dim != unpackedTypeAfterFold.getDimSize(pos))
+ return false;
+ }
return true;
}
diff --git a/mlir/test/Dialect/Linalg/canonicalize.mlir b/mlir/test/Dialect/Linalg/canonicalize.mlir
index 08eb40d4cb442..d280df614f8fc 100644
--- a/mlir/test/Dialect/Linalg/canonicalize.mlir
+++ b/mlir/test/Dialect/Linalg/canonicalize.mlir
@@ -2061,6 +2061,50 @@ func.func @no_fold_extract_slice_into_unpack_non_zero_offset(
// -----
+func.func @no_fold_extract_slice_into_unpack_slice_over_non_tiled_dim(
+ %src : tensor<30x2x16xf32>, %dest : tensor<30x32xf32>
+) -> tensor<28x28xf32> {
+ %unpack = linalg.unpack %src
+ inner_dims_pos = [1]
+ inner_tiles = [16]
+ into %dest : tensor<30x2x16xf32> -> tensor<30x32xf32>
+ %extracted_slice = tensor.extract_slice %unpack
+ [0, 0] [28, 28] [1, 1] : tensor<30x32xf32> to tensor<28x28xf32>
+ return %extracted_slice : tensor<28x28xf32>
+}
+
+// CHECK-LABEL: func @no_fold_extract_slice_into_unpack_slice_over_non_tiled_dim
+// CHECK-SAME: %[[SRC:.+]]: tensor<30x2x16xf32>
+// CHECK-SAME: %[[DEST:.+]]: tensor<30x32xf32>
+// CHECK: %[[UNPACK:.+]] = linalg.unpack %[[SRC]]
+// CHECK-SAME: into %[[DEST]]
+// CHECK: %[[SLICE:.+]] = tensor.extract_slice %[[UNPACK]]
+// CHECK: return %[[SLICE]]
+
+// -----
+
+func.func @no_fold_extract_slice_into_unpack_slice_over_dynamic_dim(
+ %src : tensor<?x2x16xf32>, %dest : tensor<?x32xf32>, %size : index
+) -> tensor<?x28xf32> {
+ %unpack = linalg.unpack %src
+ inner_dims_pos = [1]
+ inner_tiles = [16]
+ into %dest : tensor<?x2x16xf32> -> tensor<?x32xf32>
+ %extracted_slice = tensor.extract_slice %unpack
+ [0, 0] [%size, 28] [1, 1] : tensor<?x32xf32> to tensor<?x28xf32>
+ return %extracted_slice : tensor<?x28xf32>
+}
+
+// CHECK-LABEL: func @no_fold_extract_slice_into_unpack_slice_over_dynamic_dim
+// CHECK-SAME: %[[SRC:.+]]: tensor<?x2x16xf32>
+// CHECK-SAME: %[[DEST:.+]]: tensor<?x32xf32>
+// CHECK: %[[UNPACK:.+]] = linalg.unpack %[[SRC]]
+// CHECK-SAME: into %[[DEST]]
+// CHECK: %[[SLICE:.+]] = tensor.extract_slice %[[UNPACK]]
+// CHECK: return %[[SLICE]]
+
+// -----
+
// CHECK-LABEL: func.func @fold_cast_unpack_dynamic_tile_size(
// CHECK-SAME: %[[SRC:.*]]: tensor<1x1x8x1xi32>,
// CHECK-SAME: %[[DEST:.*]]: tensor<7x?xi32>) -> tensor<7x?xi32> {
>From 9544c5e45169ca61747c8d6f4f8f6c1219f5ebf6 Mon Sep 17 00:00:00 2001
From: Mikhail Romanov <mmromanov at ispras.ru>
Date: Thu, 19 Feb 2026 23:05:31 +0300
Subject: [PATCH 2/2] Add comments
---
mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp | 1 +
mlir/test/Dialect/Linalg/canonicalize.mlir | 2 ++
2 files changed, 3 insertions(+)
diff --git a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
index b02fe5564393a..ee6679f0a2e4f 100644
--- a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
@@ -6409,6 +6409,7 @@ bool UnPackOp::canFoldSliceOp(tensor::ExtractSliceOp sliceOp) {
if (paddingSize >= tileSize)
return false;
}
+ // extract_slice must not affect dimensions that are not being unpacked
for (int64_t pos = 0, e = outerShapeWithoutTranspose.size(); pos < e; ++pos) {
if (areOuterDimsTiled[pos])
continue;
diff --git a/mlir/test/Dialect/Linalg/canonicalize.mlir b/mlir/test/Dialect/Linalg/canonicalize.mlir
index d280df614f8fc..77c1c3da17166 100644
--- a/mlir/test/Dialect/Linalg/canonicalize.mlir
+++ b/mlir/test/Dialect/Linalg/canonicalize.mlir
@@ -2061,6 +2061,7 @@ func.func @no_fold_extract_slice_into_unpack_non_zero_offset(
// -----
+// Must not fold because extract_slice cuts the 0'th dimension from 30 to 28.
func.func @no_fold_extract_slice_into_unpack_slice_over_non_tiled_dim(
%src : tensor<30x2x16xf32>, %dest : tensor<30x32xf32>
) -> tensor<28x28xf32> {
@@ -2083,6 +2084,7 @@ func.func @no_fold_extract_slice_into_unpack_slice_over_non_tiled_dim(
// -----
+// Must not fold because extract_slice's effect on the 0'th dimension is unknown.
func.func @no_fold_extract_slice_into_unpack_slice_over_dynamic_dim(
%src : tensor<?x2x16xf32>, %dest : tensor<?x32xf32>, %size : index
) -> tensor<?x28xf32> {
More information about the Mlir-commits
mailing list