[Mlir-commits] [mlir] [MLIR][Vector] Generalize broadcast lowering for single-element vector to nD (PR #206501)
Artem Kroviakov
llvmlistbot at llvm.org
Tue Jun 30 02:10:44 PDT 2026
https://github.com/akroviakov updated https://github.com/llvm/llvm-project/pull/206501
>From 29026889aebcb0fadb8766d0c70983eef5f7d04c Mon Sep 17 00:00:00 2001
From: Artem Kroviakov <artem.kroviakov at intel.com>
Date: Mon, 29 Jun 2026 14:05:40 +0000
Subject: [PATCH 1/2] [MLIR][Vector] Generalize broadcast lowering for unit
vector to nD
---
.../Transforms/LowerVectorBroadcast.cpp | 6 ++--
.../VectorToLLVM/vector-to-llvm.mlir | 23 ++++++------
.../vector-broadcast-lowering-transforms.mlir | 36 +++++++++++++++++++
3 files changed, 50 insertions(+), 15 deletions(-)
diff --git a/mlir/lib/Dialect/Vector/Transforms/LowerVectorBroadcast.cpp b/mlir/lib/Dialect/Vector/Transforms/LowerVectorBroadcast.cpp
index 66dd7c8f36e6b..d2da24b366e83 100644
--- a/mlir/lib/Dialect/Vector/Transforms/LowerVectorBroadcast.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/LowerVectorBroadcast.cpp
@@ -52,9 +52,9 @@ class BroadcastOpLowering : public OpRewritePattern<vector::BroadcastOp> {
int64_t srcRank = srcType.getRank();
int64_t dstRank = dstType.getRank();
- // Here we are broadcasting to a rank-1 vector. Ensure that the source is a
- // scalar.
- if (srcRank <= 1 && dstRank == 1) {
+ // A broadcast from a single-element source vector is equivalent to scalar
+ // broadcasting to an nD shape.
+ if (srcType.getNumElements() == 1 && !srcType.isScalable()) {
SmallVector<int64_t> fullRankPosition(srcRank, 0);
Value ext = vector::ExtractOp::create(rewriter, loc, op.getSource(),
fullRankPosition);
diff --git a/mlir/test/Conversion/VectorToLLVM/vector-to-llvm.mlir b/mlir/test/Conversion/VectorToLLVM/vector-to-llvm.mlir
index 77f60b3172296..f2e9b606acd4d 100644
--- a/mlir/test/Conversion/VectorToLLVM/vector-to-llvm.mlir
+++ b/mlir/test/Conversion/VectorToLLVM/vector-to-llvm.mlir
@@ -197,17 +197,16 @@ func.func @broadcast_vec2d_from_vec0d(%arg0: vector<f32>) -> vector<3x2xf32> {
// CHECK-LABEL: @broadcast_vec2d_from_vec0d(
// CHECK-SAME: %[[A:.*]]: vector<f32>)
// CHECK: %[[T0:.*]] = builtin.unrealized_conversion_cast %[[A]] : vector<f32> to vector<1xf32>
-// CHECK: %[[T1:.*]] = ub.poison : vector<3x2xf32>
-// CHECK: %[[T2:.*]] = builtin.unrealized_conversion_cast %[[T1]] : vector<3x2xf32> to !llvm.array<3 x vector<2xf32>>
-// CHECK: %[[T4:.*]] = llvm.mlir.constant(0 : i64) : i64
-// CHECK: %[[T5:.*]] = llvm.extractelement %[[T0]][%[[T4]] : i64] : vector<1xf32>
-// CHECK: %[[T6Insert:.*]] = llvm.insertelement %[[T5]]
-// CHECK: %[[T6:.*]] = llvm.shufflevector %[[T6Insert]]
-// CHECK: %[[T7:.*]] = llvm.insertvalue %[[T6]], %[[T2]][0] : !llvm.array<3 x vector<2xf32>>
-// CHECK: %[[T8:.*]] = llvm.insertvalue %[[T6]], %[[T7]][1] : !llvm.array<3 x vector<2xf32>>
-// CHECK: %[[T9:.*]] = llvm.insertvalue %[[T6]], %[[T8]][2] : !llvm.array<3 x vector<2xf32>>
-// CHECK: %[[T10:.*]] = builtin.unrealized_conversion_cast %[[T9]] : !llvm.array<3 x vector<2xf32>> to vector<3x2xf32>
-// CHECK: return %[[T10]] : vector<3x2xf32>
+// CHECK: %[[T1:.*]] = llvm.mlir.constant(0 : i64) : i64
+// CHECK: %[[T2:.*]] = llvm.extractelement %[[T0]][%[[T1]] : i64] : vector<1xf32>
+// CHECK: %[[T3:.*]] = llvm.mlir.poison : !llvm.array<3 x vector<2xf32>>
+// CHECK: %[[T4Insert:.*]] = llvm.insertelement %[[T2]]
+// CHECK: %[[T4:.*]] = llvm.shufflevector %[[T4Insert]]
+// CHECK: %[[T5:.*]] = llvm.insertvalue %[[T4]], %[[T3]][0] : !llvm.array<3 x vector<2xf32>>
+// CHECK: %[[T6:.*]] = llvm.insertvalue %[[T4]], %[[T5]][1] : !llvm.array<3 x vector<2xf32>>
+// CHECK: %[[T7:.*]] = llvm.insertvalue %[[T4]], %[[T6]][2] : !llvm.array<3 x vector<2xf32>>
+// CHECK: %[[T8:.*]] = builtin.unrealized_conversion_cast %[[T7]] : !llvm.array<3 x vector<2xf32>> to vector<3x2xf32>
+// CHECK: return %[[T8]] : vector<3x2xf32>
// -----
@@ -1499,7 +1498,7 @@ func.func @constant_mask_2d() -> vector<4x4xi1> {
}
// CHECK-LABEL: func @constant_mask_2d
-// CHECK: %[[VAL_0:.*]] = arith.constant
+// CHECK: %[[VAL_0:.*]] = arith.constant
// CHECK-SAME{LITERAL}: dense<[[true, true, false, false], [true, true, false, false], [false, false, false, false], [false, false, false, false]]> : vector<4x4xi1>
// CHECK: return %[[VAL_0]] : vector<4x4xi1>
diff --git a/mlir/test/Dialect/Vector/vector-broadcast-lowering-transforms.mlir b/mlir/test/Dialect/Vector/vector-broadcast-lowering-transforms.mlir
index d5e344393b217..8de08e7e26417 100644
--- a/mlir/test/Dialect/Vector/vector-broadcast-lowering-transforms.mlir
+++ b/mlir/test/Dialect/Vector/vector-broadcast-lowering-transforms.mlir
@@ -95,6 +95,42 @@ func.func @broadcast_stretch(%arg0: vector<1xf32>) -> vector<4xf32> {
return %0 : vector<4xf32>
}
+// CHECK-LABEL: func @broadcast_single_elem_vec1d_to_vec2d
+// CHECK-SAME: %[[A:.*0]]: vector<1xf32>
+// CHECK: %[[T0:.*]] = vector.extract %[[A]][0] : f32 from vector<1xf32>
+// CHECK: %[[T1:.*]] = vector.broadcast %[[T0]] : f32 to vector<16x1xf32>
+// CHECK: return %[[T1]] : vector<16x1xf32>
+
+func.func @broadcast_single_elem_vec1d_to_vec2d(%arg0: vector<1xf32>) -> vector<16x1xf32> {
+ %0 = vector.broadcast %arg0 : vector<1xf32> to vector<16x1xf32>
+ return %0 : vector<16x1xf32>
+}
+
+// CHECK-LABEL: func @broadcast_single_elem_vec2d_to_vec2d
+// CHECK-SAME: %[[A:.*0]]: vector<1x1xf32>
+// CHECK: %[[T0:.*]] = vector.extract %[[A]][0, 0] : f32 from vector<1x1xf32>
+// CHECK: %[[T1:.*]] = vector.broadcast %[[T0]] : f32 to vector<4x3xf32>
+// CHECK: return %[[T1]] : vector<4x3xf32>
+
+func.func @broadcast_single_elem_vec2d_to_vec2d(%arg0: vector<1x1xf32>) -> vector<4x3xf32> {
+ %0 = vector.broadcast %arg0 : vector<1x1xf32> to vector<4x3xf32>
+ return %0 : vector<4x3xf32>
+}
+
+// CHECK-LABEL: func @broadcast_single_elem_scalable_src
+// CHECK-SAME: %[[A:.*0]]: vector<[1]xf32>
+// CHECK: %[[U0:.*]] = ub.poison : vector<4x[1]xf32>
+// CHECK: %[[T0:.*]] = vector.insert %[[A]], %[[U0]] [0] : vector<[1]xf32> into vector<4x[1]xf32>
+// CHECK: %[[T1:.*]] = vector.insert %[[A]], %[[T0]] [1] : vector<[1]xf32> into vector<4x[1]xf32>
+// CHECK: %[[T2:.*]] = vector.insert %[[A]], %[[T1]] [2] : vector<[1]xf32> into vector<4x[1]xf32>
+// CHECK: %[[T3:.*]] = vector.insert %[[A]], %[[T2]] [3] : vector<[1]xf32> into vector<4x[1]xf32>
+// CHECK: return %[[T3]] : vector<4x[1]xf32>
+
+func.func @broadcast_single_elem_scalable_src(%arg0: vector<[1]xf32>) -> vector<4x[1]xf32> {
+ %0 = vector.broadcast %arg0 : vector<[1]xf32> to vector<4x[1]xf32>
+ return %0 : vector<4x[1]xf32>
+}
+
// CHECK-LABEL: func @broadcast_stretch_at_start
// CHECK-SAME: %[[A:.*0]]: vector<1x4xf32>
// CHECK: %[[U0:.*]] = ub.poison : vector<3x4xf32>
>From 6d225601f0c576c329df64ddc83bfb8e272f11ab Mon Sep 17 00:00:00 2001
From: Artem Kroviakov <artem.kroviakov at intel.com>
Date: Tue, 30 Jun 2026 09:10:25 +0000
Subject: [PATCH 2/2] Shorter comment
---
mlir/lib/Dialect/Vector/Transforms/LowerVectorBroadcast.cpp | 3 +--
1 file changed, 1 insertion(+), 2 deletions(-)
diff --git a/mlir/lib/Dialect/Vector/Transforms/LowerVectorBroadcast.cpp b/mlir/lib/Dialect/Vector/Transforms/LowerVectorBroadcast.cpp
index d2da24b366e83..a78c31367ee28 100644
--- a/mlir/lib/Dialect/Vector/Transforms/LowerVectorBroadcast.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/LowerVectorBroadcast.cpp
@@ -52,8 +52,7 @@ class BroadcastOpLowering : public OpRewritePattern<vector::BroadcastOp> {
int64_t srcRank = srcType.getRank();
int64_t dstRank = dstType.getRank();
- // A broadcast from a single-element source vector is equivalent to scalar
- // broadcasting to an nD shape.
+ // Single-element fixed-size source: extract the scalar and broadcast it.
if (srcType.getNumElements() == 1 && !srcType.isScalable()) {
SmallVector<int64_t> fullRankPosition(srcRank, 0);
Value ext = vector::ExtractOp::create(rewriter, loc, op.getSource(),
More information about the Mlir-commits
mailing list