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

lorenzo chelini llvmlistbot at llvm.org
Thu May 16 11:00:15 PDT 2024


================
@@ -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()))
----------------
chelini wrote:

why do you want to control folding for single use? 

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


More information about the Mlir-commits mailing list