[Mlir-commits] [mlir] [mlir][Tensor] Add pattern to fold concats of empty. (PR #98994)

[llvmlistbot at llvm.org]

================
@@ -136,6 +136,38 @@ struct FoldEmptyTensorWithUnPackOp : public OpRewritePattern<UnPackOp> {
   }
 };
 
+// Fold concat operation where all the operands are empty.
+struct FoldConcatsOfEmpty : public OpRewritePattern<ConcatOp> {
+  using OpRewritePattern<ConcatOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(tensor::ConcatOp concatOp,
+                                PatternRewriter &rewriter) const override {
+    auto concatOperands = concatOp.getInputs();
+    if (concatOperands.empty()) {
+      return failure();
+    }
+    auto firstEmptyOp = concatOperands.front().getDefiningOp<tensor::EmptyOp>();
+    if (!firstEmptyOp) {
+      return failure();
+    }
+    auto isDefinedByEmptyOp = [](Value v) -> bool {
+      return v.getDefiningOp<tensor::EmptyOp>();
+    };
+    if (!llvm::all_of(concatOperands.drop_front(), isDefinedByEmptyOp)) {
+      return rewriter.notifyMatchFailure(
+          concatOp, "not all operands are defined by an empty op");
+    }
----------------
MaheshRavishankar wrote:

I was just trying to avoid even doing the `all_of` at all if even the first operand is not empty. But ok, I'll change it.

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