[Mlir-commits] [mlir] [MLIR][Linalg] Fix crash on duplicate dimensions in linalg.broadcast (PR #211203)
Chibuoyim Ogbonna
llvmlistbot at llvm.org
Wed Jul 22 05:10:50 PDT 2026
https://github.com/bruteforceboy updated https://github.com/llvm/llvm-project/pull/211203
>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 1/2] [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
>From 49929ad199e0a16f4be874624a6ec1a19f80ba1e Mon Sep 17 00:00:00 2001
From: workwilson <ogbonnachibuoyim12 at gmail.com>
Date: Wed, 22 Jul 2026 20:10:14 +0800
Subject: [PATCH 2/2] Apply review comments: trim tests, add scalar input test,
follow test naming convention
---
mlir/test/Dialect/Linalg/invalid.mlir | 27 +++++--------------
mlir/test/Dialect/Linalg/roundtrip.mlir | 36 +++----------------------
2 files changed, 10 insertions(+), 53 deletions(-)
diff --git a/mlir/test/Dialect/Linalg/invalid.mlir b/mlir/test/Dialect/Linalg/invalid.mlir
index 45bd8b87c0970..8c3f27efcd58a 100644
--- a/mlir/test/Dialect/Linalg/invalid.mlir
+++ b/mlir/test/Dialect/Linalg/invalid.mlir
@@ -1208,27 +1208,14 @@ func.func @broadcast_duplicate_dims(
// -----
-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}}
+func.func @broadcast_scalar_input(
+ %input: f32, %init: tensor<8x16xf32>) -> tensor<8x16xf32> {
+ // expected-error @+1 {{'linalg.broadcast' op operand #0 must be memref of any non-token type values or ranked tensor of any non-token type values, but got 'f32'}}
%bcast = linalg.broadcast
- ins(%input:tensor<4x8xf32>)
- outs(%init:tensor<4x8x16x32xf32>)
- dimensions = [3, 3]
- func.return %bcast : tensor<4x8x16x32xf32>
+ ins(%input:f32)
+ outs(%init:tensor<8x16xf32>)
+ dimensions = [0, 1]
+ func.return %bcast : tensor<8x16xf32>
}
// -----
diff --git a/mlir/test/Dialect/Linalg/roundtrip.mlir b/mlir/test/Dialect/Linalg/roundtrip.mlir
index f2584f7b2ecec..5f9e98bdadb5d 100644
--- a/mlir/test/Dialect/Linalg/roundtrip.mlir
+++ b/mlir/test/Dialect/Linalg/roundtrip.mlir
@@ -599,51 +599,21 @@ func.func @broadcast_memref(%input: memref<8x32xf32>,
// -----
-func.func @broadcast_0d_to_2d(%input: tensor<i32>,
- %init: tensor<32x2xi32>) -> tensor<32x2xi32> {
+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_to_2d
+// CHECK-LABEL: func @broadcast_0d_tensor
// 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
More information about the Mlir-commits
mailing list