[Mlir-commits] [mlir] f92e977 - [MLIR][Linalg] Fix crash on duplicate dimensions in linalg.broadcast (#211203)

[llvmlistbot at llvm.org]

Author: Chibuoyim (Wilson) Ogbonna
Date: 2026-07-27T13:02:04+01:00
New Revision: f92e977a8e76a0e40e3af255f13c30b6710fd871

URL: https://github.com/llvm/llvm-project/commit/f92e977a8e76a0e40e3af255f13c30b6710fd871
DIFF: https://github.com/llvm/llvm-project/commit/f92e977a8e76a0e40e3af255f13c30b6710fd871.diff

LOG: [MLIR][Linalg] Fix crash on duplicate dimensions in linalg.broadcast (#211203)

per
[#210738](https://github.com/llvm/llvm-project/pull/210738#discussion_r3622744153)

Currently, the following crashes on verification with a segfault:
```
func.func @broadcast_rank0_tensor(%arg0: tensor<i32>, %arg1: tensor<32x2xi32>) -> tensor<32x2xi32> {
  %0 = linalg.broadcast ins(%arg0 : tensor<i32>) outs(%arg1 : tensor<32x2xi32>) dimensions = [0, 0]
  return %0 : tensor<32x2xi32>
}
```

The verifier should reject cases like this where we have a duplicate
dimension. I've also added a couple of tests covering `linalg.broadcast`
on various input scenarios to confirm they are being handled properly.

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td
    mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
    mlir/test/Dialect/Linalg/invalid.mlir
    mlir/test/Dialect/Linalg/roundtrip.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td b/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td
index fc5c9770d969b..fa8125f280db2 100644
--- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td
+++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td
@@ -478,6 +478,11 @@ def BroadcastOp : LinalgStructuredBase_Op<"broadcast", [
   let description = [{
     Broadcast the input into the given shape by adding `dimensions`.
 
+    Each index in the `dimensions` attribute refers to a dimension of `init`
+    that is added by the operation. The indices must be unique and within the
+    rank of `init`; the sizes of the remaining (non-added) dimensions of
+    `init` must match the shape of `input`.
+
     Example:
     ```mlir
       %bcast = linalg.broadcast

diff  --git a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
index 8b2f5064e78bd..170e1edf8a55d 100644
--- a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
@@ -2356,6 +2356,10 @@ LogicalResult BroadcastOp::verify() {
                            << initRank - 1 << "], got: " << dim;
   }
 
+  DenseSet<int64_t> uniquedDims(llvm::from_range, dimensionsRef);
+  if (uniquedDims.size() != dimensionsRef.size())
+    return emitOpError() << "dimensions should not contain duplicates";
+
   // Mapping from input dims to init dims.
   SmallVector<int64_t> dimMap;
   for (auto dim : llvm::seq<int64_t>(0, initRank)) {

diff  --git a/mlir/test/Dialect/Linalg/invalid.mlir b/mlir/test/Dialect/Linalg/invalid.mlir
index c52163e244ea2..a5ec1fcc58880 100644
--- a/mlir/test/Dialect/Linalg/invalid.mlir
+++ b/mlir/test/Dialect/Linalg/invalid.mlir
@@ -1196,6 +1196,18 @@ func.func @broadcast_size_1_extension_not_supported(
 
 // -----
 
+func.func @broadcast_duplicate_dims(
+    %input: tensor<i32>, %init: tensor<32x2xi32>) -> tensor<32x2xi32> {
+  // expected-error @+1 {{'linalg.broadcast' op dimensions should not contain duplicates}}
+  %bcast = linalg.broadcast
+      ins(%input:tensor<i32>)
+      outs(%init:tensor<32x2xi32>)
+      dimensions = [0, 0]
+  func.return %bcast : tensor<32x2xi32>
+}
+
+// -----
+
 func.func @broadcast_no_operands1() {
   // expected-error @+1 {{'linalg.broadcast' op expected 2 operands, but found 0}}
   linalg.broadcast dimensions = [1]

diff  --git a/mlir/test/Dialect/Linalg/roundtrip.mlir b/mlir/test/Dialect/Linalg/roundtrip.mlir
index bfb92c3289a49..2ad69d654ab99 100644
--- a/mlir/test/Dialect/Linalg/roundtrip.mlir
+++ b/mlir/test/Dialect/Linalg/roundtrip.mlir
@@ -599,6 +599,21 @@ func.func @broadcast_memref(%input: memref<8x32xf32>,
 
 // -----
 
+func.func @broadcast_0d_tensor(%input: tensor<i32>,
+                               %init: tensor<32x2xi32>) -> tensor<32x2xi32> {
+  %bcast = linalg.broadcast
+      ins(%input:tensor<i32>)
+      outs(%init:tensor<32x2xi32>)
+      dimensions = [0, 1]
+  func.return %bcast : tensor<32x2xi32>
+}
+// CHECK-LABEL: func @broadcast_0d_tensor
+//      CHECK:    linalg.broadcast ins(%{{.*}} : tensor<i32>)
+// CHECK-SAME:    outs(%{{.*}} : tensor<32x2xi32>)
+// CHECK-SAME:    dimensions = [0, 1]
+
+// -----
+
 func.func @map_arith_with_attr(%lhs: tensor<64xf32>, %rhs: tensor<64xf32>,
                       %init: tensor<64xf32>) -> tensor<64xf32> {
   %add = linalg.map