[Mlir-commits] [mlir] [mlir][tensor] Add FoldExtractSliceOfExpandShape folding pattern (PR #212974)

[Tuomas Kärnä]

https://github.com/tkarna updated https://github.com/llvm/llvm-project/pull/212974

>From eefad915a893a22ccb2bba59806470c1951c0f0c 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/2] tensor: add FoldExtractSliceOfExpandShape to
 ReassociativeReshapeFoldingPatterns

---
 .../Tensor/Transforms/ReshapePatterns.cpp     | 48 ++++++++++++++++---
 .../Tensor/fold-reassociative-reshapes.mlir   | 31 ++++++++++++
 2 files changed, 72 insertions(+), 7 deletions(-)

diff --git a/mlir/lib/Dialect/Tensor/Transforms/ReshapePatterns.cpp b/mlir/lib/Dialect/Tensor/Transforms/ReshapePatterns.cpp
index 815ddaf630b2a..e2d541c6f6b0b 100644
--- a/mlir/lib/Dialect/Tensor/Transforms/ReshapePatterns.cpp
+++ b/mlir/lib/Dialect/Tensor/Transforms/ReshapePatterns.cpp
@@ -52,6 +52,40 @@ struct FoldExpandOfRankReducingExtract
   }
 };
 
+/// 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.
+struct FoldExtractSliceOfExpandShape : public OpRewritePattern<ExtractSliceOp> {
+  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 collapse_shape which only removes static dimensions of size `1`
 /// into extract_slice.
 struct FoldUnPaddingCollapseIntoExtract
@@ -808,13 +842,13 @@ LogicalResult mlir::tensor::getExpandedExtractSliceInfo(
 
 void mlir::tensor::populateReassociativeReshapeFoldingPatterns(
     RewritePatternSet &patterns) {
-  patterns
-      .add<FoldExpandOfRankReducingExtract, FoldUnPaddingCollapseIntoExtract,
-           FoldInsertOfRankReducingInsert<tensor::InsertSliceOp>,
-           FoldInsertOfRankReducingInsert<tensor::ParallelInsertSliceOp>,
-           FoldPaddingExpandIntoInsert<tensor::InsertSliceOp>,
-           FoldPaddingExpandIntoInsert<tensor::ParallelInsertSliceOp>>(
-          patterns.getContext());
+  patterns.add<FoldExpandOfRankReducingExtract, FoldExtractSliceOfExpandShape,
+               FoldUnPaddingCollapseIntoExtract,
+               FoldInsertOfRankReducingInsert<tensor::InsertSliceOp>,
+               FoldInsertOfRankReducingInsert<tensor::ParallelInsertSliceOp>,
+               FoldPaddingExpandIntoInsert<tensor::InsertSliceOp>,
+               FoldPaddingExpandIntoInsert<tensor::ParallelInsertSliceOp>>(
+      patterns.getContext());
 }
 
 void mlir::tensor::populateBubbleUpExpandShapePatterns(
diff --git a/mlir/test/Dialect/Tensor/fold-reassociative-reshapes.mlir b/mlir/test/Dialect/Tensor/fold-reassociative-reshapes.mlir
index 594d540dfca0a..d9ff20deb034f 100644
--- a/mlir/test/Dialect/Tensor/fold-reassociative-reshapes.mlir
+++ b/mlir/test/Dialect/Tensor/fold-reassociative-reshapes.mlir
@@ -24,6 +24,36 @@ func.func @expand_shape_of_rank_reducing_extract(
 
 // -----
 
+// 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 @unpadding_collapse_of_extract_slice(
 //  CHECK-SAME:     %[[t:.*]]: tensor<?x?x?x?xf32>
 //  CHECK-SAME:     %[[x:[a-zA-Z0-9_]+]]: index
@@ -238,3 +268,4 @@ func.func @parallel_insert_of_non_padding_expand_shape(
   }
   return %1 : tensor<?x?x?x?xf32>
 }
+

>From f505432ed9d2d7b0471dbcab60577385893d9544 Mon Sep 17 00:00:00 2001
From: Tuomas Karna <tuomas.karna at intel.com>
Date: Mon, 24 Aug 2026 11:11:33 +0300
Subject: [PATCH 2/2] add notifyMatchFailures

---
 .../Dialect/Tensor/Transforms/ReshapePatterns.cpp    | 12 ++++++++----
 1 file changed, 8 insertions(+), 4 deletions(-)

diff --git a/mlir/lib/Dialect/Tensor/Transforms/ReshapePatterns.cpp b/mlir/lib/Dialect/Tensor/Transforms/ReshapePatterns.cpp
index e2d541c6f6b0b..c9bd4181da4f8 100644
--- a/mlir/lib/Dialect/Tensor/Transforms/ReshapePatterns.cpp
+++ b/mlir/lib/Dialect/Tensor/Transforms/ReshapePatterns.cpp
@@ -64,21 +64,25 @@ struct FoldExtractSliceOfExpandShape : public OpRewritePattern<ExtractSliceOp> {
       return failure();
 
     if (sliceOp.getType() != expandOp.getSrcType())
-      return failure();
+      return rewriter.notifyMatchFailure(
+          sliceOp, "slice result type does not match expand_shape source type");
 
     SmallVector<OpFoldResult> mixedExpandedSizes =
         expandOp.getMixedOutputShape();
     if (mixedExpandedSizes.size() != sliceOp.getMixedSizes().size())
-      return failure();
+      return rewriter.notifyMatchFailure(
+          sliceOp, "expand_shape output rank does not match slice rank");
 
     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();
+        return rewriter.notifyMatchFailure(
+            sliceOp, "slice is not a zero-offset, unit-stride full slice");
       if (size != expandedSize)
-        return failure();
+        return rewriter.notifyMatchFailure(
+            sliceOp, "slice size does not match expand_shape output size");
     }
 
     rewriter.replaceOp(sliceOp, expandOp.getSrc());