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

[Adam Siemieniuk]

================
@@ -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();
----------------
adam-smnk wrote:

nit: I'd consider adding some `notifyMatchFailure` reasons for more context but it's optional

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