[Mlir-commits] [mlir] [mlir][tensor] Add more tensor.extract_slice canonicalization (PR #212974)

[Tuomas Kärnä]

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

>From 47f943c49627ff582ea18de8c97560a4a7ba88ac 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] tensor: add FoldExtractSliceOfExpandShape to
 ReassociativeReshapeFoldingPatterns

---
 .../Tensor/Transforms/ReshapePatterns.cpp     | 37 ++++++++++++++++++-
 .../Tensor/fold-reassociative-reshapes.mlir   | 31 ++++++++++++++++
 2 files changed, 67 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/Tensor/Transforms/ReshapePatterns.cpp b/mlir/lib/Dialect/Tensor/Transforms/ReshapePatterns.cpp
index 815ddaf630b2a..bb8444e813db5 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
@@ -809,7 +843,8 @@ LogicalResult mlir::tensor::getExpandedExtractSliceInfo(
 void mlir::tensor::populateReassociativeReshapeFoldingPatterns(
     RewritePatternSet &patterns) {
   patterns
-      .add<FoldExpandOfRankReducingExtract, FoldUnPaddingCollapseIntoExtract,
+      .add<FoldExpandOfRankReducingExtract, FoldExtractSliceOfExpandShape,
+           FoldUnPaddingCollapseIntoExtract,
            FoldInsertOfRankReducingInsert<tensor::InsertSliceOp>,
            FoldInsertOfRankReducingInsert<tensor::ParallelInsertSliceOp>,
            FoldPaddingExpandIntoInsert<tensor::InsertSliceOp>,
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>
 }
+