[Mlir-commits] [mlir] [mlir][Linalg] Allow isaBroadcastOpInterface to accept LinalgOp (PR #182806)

[Abhishek Varma]

https://github.com/Abhishek-Varma created https://github.com/llvm/llvm-project/pull/182806

Allow isaBroadcastOpInterface to accept LinalgOp so that both the named linalg.broadcast op and broadcast-like linalg.generic are handled by a single API instead of special-casing check for named vs generic op in downstream projects.

No test is being added for this change because callers that pass GenericOp (e.g. Specialize.cpp) continue to work since GenericOp is a LinalgOp.

Signed-off-by: Abhishek Varma <abhvarma at amd.com>

>From 9f23262e539c4d9d61a2924af481786d22ce856d Mon Sep 17 00:00:00 2001
From: IREE <no-reply at iree.dev>
Date: Mon, 23 Feb 2026 05:32:27 +0000
Subject: [PATCH] [mlir][Linalg] Allow isaBroadcastOpInterface to accept
 LinalgOp

Allow isaBroadcastOpInterface to accept LinalgOp so that both the named
linalg.broadcast op and broadcast-like linalg.generic are handled by a
single API instead of special-casing check for named vs generic op in
downstream projects.

No test is being added for this change because callers that pass
GenericOp (e.g. Specialize.cpp) continue to work since GenericOp is
a LinalgOp.

Signed-off-by: Abhishek Varma <abhvarma at amd.com>
---
 mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.h |  7 +++----
 mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp        | 10 +++++++++-
 2 files changed, 12 insertions(+), 5 deletions(-)

diff --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.h b/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.h
index 0ebbeea937554..81d51153352da 100644
--- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.h
+++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.h
@@ -121,10 +121,9 @@ bool isaConvolutionOpInterface(LinalgOp linalgOp,
 /// Checks whether `linalgOp` is semantically equivalent to a `linalg.copyOp`.
 bool isaCopyOpInterface(LinalgOp linalgOp);
 
-/// Checks whether `genericOp` is semantically equivalent to a
-///  `linalg.broadcast`. Returns broadcast dimensions if true.
-std::optional<SmallVector<int64_t>>
-isaBroadcastOpInterface(GenericOp genericOp);
+/// Checks whether `linalgOp` is semantically equivalent to a broadcast
+/// operation. Returns broadcast dimensions if true.
+std::optional<SmallVector<int64_t>> isaBroadcastOpInterface(LinalgOp linalgOp);
 
 /// Checks whether `genericOp` is semantically equivalent to a
 ///  `linalg.transpose`. Returns permuted dimensions if true.
diff --git a/mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp
index b4b1347493529..df5e258f149d4 100644
--- a/mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp
@@ -138,7 +138,15 @@ std::optional<Value> linalg::isaFillOpInterface(GenericOp op) {
 // BroadcastOpInterface implementation
 //===----------------------------------------------------------------------===//
 std::optional<SmallVector<int64_t>>
-linalg::isaBroadcastOpInterface(GenericOp op) {
+linalg::isaBroadcastOpInterface(LinalgOp linalgOp) {
+  if (auto broadcastOp = dyn_cast<BroadcastOp>(linalgOp.getOperation()))
+    return SmallVector<int64_t>(broadcastOp.getDimensions().begin(),
+                                broadcastOp.getDimensions().end());
+
+  auto op = dyn_cast<GenericOp>(linalgOp.getOperation());
+  if (!op)
+    return std::nullopt;
+
   // Structural.
   if (!op.isAllParallelLoops() || !op.isSingleInputOutput() ||
       !op.isSingleYieldOp())