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

[llvmlistbot at llvm.org]

Author: Tuomas Kärnä
Date: 2026-08-25T17:48:31+03:00
New Revision: c7beee7925352c2fbd01e0b482d124074ee80852

URL: https://github.com/llvm/llvm-project/commit/c7beee7925352c2fbd01e0b482d124074ee80852
DIFF: https://github.com/llvm/llvm-project/commit/c7beee7925352c2fbd01e0b482d124074ee80852.diff

LOG: [mlir][tensor] Add FoldExtractSliceOfExpandShape folding pattern (#212974)

Extends Tensor `ReassociativeReshapeFoldingPatterns` by adding a
`FoldExtractSliceOfExpandShape` pattern.

FoldExtractSliceOfExpandShape pattern folds full-slice rank-reducing
extract of expand_shape.

Assisted-by: GPT-5.3-Codex

Added: 
    

Modified: 
    mlir/lib/Dialect/Tensor/Transforms/ReshapePatterns.cpp
    mlir/test/Dialect/Tensor/fold-reassociative-reshapes.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Tensor/Transforms/ReshapePatterns.cpp b/mlir/lib/Dialect/Tensor/Transforms/ReshapePatterns.cpp
index 815ddaf630b2a..c9bd4181da4f8 100644
--- a/mlir/lib/Dialect/Tensor/Transforms/ReshapePatterns.cpp
+++ b/mlir/lib/Dialect/Tensor/Transforms/ReshapePatterns.cpp
@@ -52,6 +52,44 @@ 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 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 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 rewriter.notifyMatchFailure(
+            sliceOp, "slice is not a zero-offset, unit-stride full slice");
+      if (size != expandedSize)
+        return rewriter.notifyMatchFailure(
+            sliceOp, "slice size does not match expand_shape output size");
+    }
+
+    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 +846,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_
diff erent_sizes(
+//       CHECK:   tensor.expand_shape
+//       CHECK:   tensor.extract_slice
+func.func @dont_fold_extract_slice_of_expand_shape_with_
diff erent_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>
 }
+