[Mlir-commits] [mlir] eff826e - [MLIR][Vector] Generalize broadcast lowering for single-element vector to nD (#206501)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Jul 1 08:21:41 PDT 2026
Author: Artem Kroviakov
Date: 2026-07-01T15:21:35Z
New Revision: eff826e91dc7321eb38821da600dd411707e67fa
URL: https://github.com/llvm/llvm-project/commit/eff826e91dc7321eb38821da600dd411707e67fa
DIFF: https://github.com/llvm/llvm-project/commit/eff826e91dc7321eb38821da600dd411707e67fa.diff
LOG: [MLIR][Vector] Generalize broadcast lowering for single-element vector to nD (#206501)
Added:
Modified:
mlir/lib/Dialect/Vector/Transforms/LowerVectorBroadcast.cpp
mlir/test/Conversion/VectorToLLVM/vector-to-llvm.mlir
mlir/test/Dialect/Vector/vector-broadcast-lowering-transforms.mlir
Removed:
################################################################################
diff --git a/mlir/lib/Dialect/Vector/Transforms/LowerVectorBroadcast.cpp b/mlir/lib/Dialect/Vector/Transforms/LowerVectorBroadcast.cpp
index 66dd7c8f36e6b..a78c31367ee28 100644
--- a/mlir/lib/Dialect/Vector/Transforms/LowerVectorBroadcast.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/LowerVectorBroadcast.cpp
@@ -52,9 +52,8 @@ 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) {
+ // 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(),
fullRankPosition);
diff --git a/mlir/test/Conversion/VectorToLLVM/vector-to-llvm.mlir b/mlir/test/Conversion/VectorToLLVM/vector-to-llvm.mlir
index 1fc80ffb5cc50..094353119e493 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..5b193c5e25ca0 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_vec2d_from_vec1d_single_element
+// 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_vec2d_from_vec1d_single_element(%arg0: vector<1xf32>) -> vector<16x1xf32> {
+ %0 = vector.broadcast %arg0 : vector<1xf32> to vector<16x1xf32>
+ return %0 : vector<16x1xf32>
+}
+
+// CHECK-LABEL: func @broadcast_vec2d_from_vec2d_single_element
+// 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_vec2d_from_vec2d_single_element(%arg0: vector<1x1xf32>) -> vector<4x3xf32> {
+ %0 = vector.broadcast %arg0 : vector<1x1xf32> to vector<4x3xf32>
+ return %0 : vector<4x3xf32>
+}
+
+// CHECK-LABEL: func @broadcast_vec2d_from_vec1d_unit_scalable_dim
+// 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_vec2d_from_vec1d_unit_scalable_dim(%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>
More information about the Mlir-commits
mailing list