[Mlir-commits] [mlir] [MLIR][Linalg] Fix crash on duplicate dimensions in linalg.broadcast (PR #211203)
Chibuoyim Ogbonna
llvmlistbot at llvm.org
Wed Jul 22 01:49:22 PDT 2026
https://github.com/bruteforceboy created https://github.com/llvm/llvm-project/pull/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 bug was found on 0D inputs but extends to other ranks. 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 0D input scenarios to confirm they are being handled properly.
>From 19475b38b64fc264dca3ea3dac589affafe74bbb Mon Sep 17 00:00:00 2001
From: workwilson <ogbonnachibuoyim12 at gmail.com>
Date: Wed, 22 Jul 2026 16:37:17 +0800
Subject: [PATCH] [MLIR][Linalg] Fix crash on duplicate dimensions in
linalg.broadcast
---
.../Dialect/Linalg/IR/LinalgStructuredOps.td | 5 ++
mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp | 4 ++
mlir/test/Dialect/Linalg/invalid.mlir | 37 ++++++++++++
mlir/test/Dialect/Linalg/roundtrip.mlir | 60 +++++++++++++++++++
4 files changed, 106 insertions(+)
diff --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td b/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td
index fc5c9770d969b..ddb4f6675103e 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 `dimensions` attribute refers to a dimension of the init
+ tensor that is added by the operation. The indices must be unique and
+ within the init rank; the sizes of the remaining (non-added) init
+ dimensions must match the input shape.
+
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..45bd8b87c0970 100644
--- a/mlir/test/Dialect/Linalg/invalid.mlir
+++ b/mlir/test/Dialect/Linalg/invalid.mlir
@@ -1196,6 +1196,43 @@ 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_duplicate_dims_rank1(
+ %input: tensor<4xf32>, %init: tensor<4x8x16xf32>) -> tensor<4x8x16xf32> {
+ // expected-error @+1 {{'linalg.broadcast' op dimensions should not contain duplicates}}
+ %bcast = linalg.broadcast
+ ins(%input:tensor<4xf32>)
+ outs(%init:tensor<4x8x16xf32>)
+ dimensions = [1, 1]
+ func.return %bcast : tensor<4x8x16xf32>
+}
+
+// -----
+
+func.func @broadcast_duplicate_dims_rank2(
+ %input: tensor<4x8xf32>, %init: tensor<4x8x16x32xf32>)
+ -> tensor<4x8x16x32xf32> {
+ // expected-error @+1 {{'linalg.broadcast' op dimensions should not contain duplicates}}
+ %bcast = linalg.broadcast
+ ins(%input:tensor<4x8xf32>)
+ outs(%init:tensor<4x8x16x32xf32>)
+ dimensions = [3, 3]
+ func.return %bcast : tensor<4x8x16x32xf32>
+}
+
+// -----
+
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..f2584f7b2ecec 100644
--- a/mlir/test/Dialect/Linalg/roundtrip.mlir
+++ b/mlir/test/Dialect/Linalg/roundtrip.mlir
@@ -599,6 +599,66 @@ func.func @broadcast_memref(%input: memref<8x32xf32>,
// -----
+func.func @broadcast_0d_to_2d(%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_to_2d
+// CHECK: linalg.broadcast ins(%{{.*}} : tensor<i32>)
+// CHECK-SAME: outs(%{{.*}} : tensor<32x2xi32>)
+// CHECK-SAME: dimensions = [0, 1]
+
+// -----
+
+func.func @broadcast_0d_to_1d(%input: tensor<f32>,
+ %init: tensor<64xf32>) -> tensor<64xf32> {
+ %bcast = linalg.broadcast
+ ins(%input:tensor<f32>)
+ outs(%init:tensor<64xf32>)
+ dimensions = [0]
+ func.return %bcast : tensor<64xf32>
+}
+// CHECK-LABEL: func @broadcast_0d_to_1d
+// CHECK: linalg.broadcast ins(%{{.*}} : tensor<f32>)
+// CHECK-SAME: outs(%{{.*}} : tensor<64xf32>)
+// CHECK-SAME: dimensions = [0]
+
+// -----
+
+func.func @broadcast_0d_with_dynamic_sizes(
+ %input: tensor<f32>, %init: tensor<8x?xf32>) -> tensor<8x?xf32> {
+ %bcast = linalg.broadcast
+ ins(%input:tensor<f32>)
+ outs(%init:tensor<8x?xf32>)
+ dimensions = [0, 1]
+ func.return %bcast : tensor<8x?xf32>
+}
+// CHECK-LABEL: func @broadcast_0d_with_dynamic_sizes
+// CHECK: linalg.broadcast ins(%{{.*}} : tensor<f32>)
+// CHECK-SAME: outs(%{{.*}} : tensor<8x?xf32>)
+// CHECK-SAME: dimensions = [0, 1]
+
+// -----
+
+func.func @broadcast_0d_memref(%input: memref<f32>,
+ %init: memref<8x16xf32>) {
+ linalg.broadcast
+ ins(%input:memref<f32>)
+ outs(%init:memref<8x16xf32>)
+ dimensions = [0, 1]
+ func.return
+}
+// CHECK-LABEL: func @broadcast_0d_memref
+// CHECK: linalg.broadcast ins(%{{.*}} : memref<f32>)
+// CHECK-SAME: outs(%{{.*}} : memref<8x16xf32>)
+// 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
More information about the Mlir-commits
mailing list