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

[Han-Chung Wang]

================
@@ -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");
+    }
----------------
hanhanW wrote:

I thought that the `llvm::all_of` should already consider the case? I imagine that it also does early-exit if the first operand is not an empty op. It could be different in multi-threads though, idk.

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