[Mlir-commits] [mlir] [mlir][linalg]Implement canonicalizer for `BroadCastOp` on tensors (PR #80466)

[llvmlistbot at llvm.org]

https://github.com/srcarroll created https://github.com/llvm/llvm-project/pull/80466

None

>From f51bb7b15be55e682be76f2289b991ed42ab4d41 Mon Sep 17 00:00:00 2001
From: Sam <srcarroll314 at gmail.com>
Date: Fri, 2 Feb 2024 11:37:03 -0600
Subject: [PATCH] [mlir][linalg]Implement canonicalizer for
 `linalg::BroadCastOp` on tensors

---
 .../mlir/Dialect/Linalg/IR/LinalgStructuredOps.td  |  1 +
 mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp           | 14 ++++++++++++++
 2 files changed, 15 insertions(+)

diff --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td b/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td
index 751edd0228830..11b6f50032c09 100644
--- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td
+++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td
@@ -531,6 +531,7 @@ def BroadcastOp : LinalgStructuredBase_Op<"broadcast", [
 
   let hasCustomAssemblyFormat = 1;
   let hasVerifier = 1;
+  let hasCanonicalizeMethod = 1;
 }
 
 //===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
index e86b9762d8581..cddb0671dd58f 100644
--- a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
@@ -1907,6 +1907,20 @@ void BroadcastOp::getEffects(
                         getDpsInits());
 }
 
+LogicalResult BroadcastOp::canonicalize(BroadcastOp op,
+                                        PatternRewriter &rewriter) {
+  // For tensor semantics, if op's input and init are same shape, it is a no op.
+  // Otherwise, with buffer semantics, the op does a copy and we don't
+  // canonicalize.
+  if (op.hasPureTensorSemantics() &&
+      (op.getInput().getType() == op.getInit().getType())) {
+    rewriter.replaceAllUsesWith(op.getResult(), op.getInput());
+    rewriter.eraseOp(op);
+    return success();
+  }
+  return failure();
+}
+
 //===----------------------------------------------------------------------===//
 // YieldOp
 //===----------------------------------------------------------------------===//