[Mlir-commits] [mlir] [mlir][tensor] Fold pack and unpack of empty input tensor (PR #92247)

[Adam Siemieniuk]

================
@@ -93,12 +93,76 @@ struct FoldEmptyTensorWithExtractSliceOp
   bool foldSingleUseOnly = false;
 };
 
+/// tensor.empty does not define any tensor contents, so an unpadded pack
+/// can be folded away.
+struct FoldEmptyTensorWithPackOp : public OpRewritePattern<PackOp> {
+  FoldEmptyTensorWithPackOp(MLIRContext *ctx, PatternBenefit benefit = 1,
+                            bool foldSingleUseOnly = false)
+      : OpRewritePattern<PackOp>(ctx, benefit),
+        foldSingleUseOnly(foldSingleUseOnly) {}
+
+  LogicalResult matchAndRewrite(PackOp packOp,
+                                PatternRewriter &rewriter) const override {
+    // Check for tensor.empty source.
+    auto emptyOp = packOp.getSource().getDefiningOp<EmptyOp>();
+    if (!emptyOp)
+      return failure();
+
+    // Check for single use.
+    if (foldSingleUseOnly && !llvm::hasSingleElement(emptyOp->getUses()))
----------------
adam-smnk wrote:

Dropped the check for pack/unpack.
I don't think there's any scenario where multiple tensor empty users could prevent this fold pattern.

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