[Mlir-commits] [mlir] [mlir] [vector] try promoting scalar when reordering broadcast/elementwise (PR #212180)
Chuanqi Xu
llvmlistbot at llvm.org
Sun Jul 26 23:38:16 PDT 2026
https://github.com/ChuanqiXu9 created https://github.com/llvm/llvm-project/pull/212180
Inspired by https://github.com/llvm/llvm-project/pull/211208
The thread discussed the case:
```
%0 = vector.broadcast %arg0 : vector<4xf32> to vector<3x4xf32>
%1 = vector.broadcast %arg1 : f32 to vector<3x4xf32>
%2 = vector.broadcast %arg2 : vector<4xf32> to vector<3x4xf32>
%3 = vector.fma %0, %1, %2 : vector<3x4xf32>
```
The reviewer suggests "broadcasts on %arg0 and %arg2 are removed but %arg1 is broadcasted to <4xf32>. Then FMA would happen on <4xf32>, %followed by the broadcast to <3x4xf32>". Now this is solved and we can see the test case at @fma_mixed_scalar_and_vector_broadcast_source in the attached test case.
AI assisted.
>From 68ec23f2dd692ac2e07a96521b1c9b3a64c64b61 Mon Sep 17 00:00:00 2001
From: "yedeng.yd" <yedeng.yd at alibaba-inc.com>
Date: Mon, 27 Jul 2026 14:30:36 +0800
Subject: [PATCH] [mlir] [vector] try promoting scalar when reordering
broadcast/elementwise
Inspired by https://github.com/llvm/llvm-project/pull/211208
The thread discussed the case:
```
%0 = vector.broadcast %arg0 : vector<4xf32> to vector<3x4xf32>
%1 = vector.broadcast %arg1 : f32 to vector<3x4xf32>
%2 = vector.broadcast %arg2 : vector<4xf32> to vector<3x4xf32>
%3 = vector.fma %0, %1, %2 : vector<3x4xf32>
```
The reviewer suggests "broadcasts on %arg0 and %arg2 are removed but
%arg1 is broadcasted to <4xf32>. Then FMA would happen on <4xf32>,
%followed by the broadcast to <3x4xf32>". Now this is solved and we
can see the test case at @fma_mixed_scalar_and_vector_broadcast_source
in the attached test case.
AI assisted.
---
.../Vector/Transforms/VectorTransforms.cpp | 48 ++++++++++----
mlir/test/Dialect/Vector/vector-sink.mlir | 65 ++++++++++++++-----
2 files changed, 83 insertions(+), 30 deletions(-)
diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
index d7beb08d3333c..b2626d17b8bbd 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
@@ -1046,41 +1046,54 @@ struct ReorderElementwiseOpsOnBroadcast final
Type resultElemType = resultType.getElementType();
- // Get the type of the first non-constant operand
+ // Prefer a vector source when one is available. This is important for
+ // mixed scalar/vector broadcasts: the vector source determines the shape
+ // to which the scalar is promoted.
Value broadcastSource;
+ Value firstBroadcastSource;
for (Value operand : op->getOperands()) {
Operation *definingOp = operand.getDefiningOp();
if (!definingOp)
return failure();
if (definingOp->hasTrait<OpTrait::ConstantLike>())
continue;
- broadcastSource = getBroadcastLikeSource(operand);
- break;
+ Value source = getBroadcastLikeSource(operand);
+ if (!source)
+ return failure();
+ if (!firstBroadcastSource)
+ firstBroadcastSource = source;
+ if (isa<VectorType>(source.getType())) {
+ broadcastSource = source;
+ break;
+ }
}
+ if (!broadcastSource)
+ broadcastSource = firstBroadcastSource;
if (!broadcastSource)
return failure();
Type unbroadcastResultType =
cloneOrReplace(broadcastSource.getType(), resultElemType);
- // Some ops, e.g. `vector.fma`, only accept vector types. For such ops the
- // reordering is only possible when the broadcast source is a vector as
- // well; sinking past a broadcast from a scalar would create an invalid op.
- // TODO: It may be better to support scalar sources by promoting the scalar
- // to a single element vector.
+ // Some ops, e.g. `vector.fma`, only accept vector types. For such ops, a
+ // vector broadcast source is needed to determine the type of the reordered
+ // op. Scalar sources can then be promoted to that vector type.
+ // TODO: Support the case where all broadcast sources are scalars by
+ // promoting them to single element vectors.
if (isa<vector::FMAOp>(op) && !isa<VectorType>(unbroadcastResultType)) {
return rewriter.notifyMatchFailure(
op, "Op only accepts vector types, but the broadcast source is a "
"scalar");
}
- // Make sure that all operands are broadcast from identically-shaped types:
- // * scalar (`vector.broadcast`), or
- // * vector (`vector.broadcast`).
- // Otherwise the re-ordering wouldn't be safe.
+ // Make sure that all operands are broadcasts from compatible source types.
+ // Scalar sources are allowed when a vector source is available and are
+ // promoted to the vector source type selected above.
if (!llvm::all_of(op->getOperands(), [broadcastSource](Value val) {
if (auto source = getBroadcastLikeSource(val))
return haveSameShapeAndScaling(source.getType(),
- broadcastSource.getType());
+ broadcastSource.getType()) ||
+ (isa<VectorType>(broadcastSource.getType()) &&
+ !isa<VectorType>(source.getType()));
SplatElementsAttr splatConst;
return matchPattern(val, m_Constant(&splatConst));
})) {
@@ -1108,7 +1121,14 @@ struct ReorderElementwiseOpsOnBroadcast final
rewriter, newConst, newType, operand.getLoc());
srcValues.push_back(newConstOp->getResult(0));
} else {
- srcValues.push_back(operand.getDefiningOp()->getOperand(0));
+ Value source = operand.getDefiningOp()->getOperand(0);
+ if (isa<VectorType>(broadcastSource.getType()) &&
+ !isa<VectorType>(source.getType()))
+ source = vector::BroadcastOp::create(
+ rewriter, operand.getLoc(),
+ cloneOrReplace(broadcastSource.getType(), source.getType()),
+ source);
+ srcValues.push_back(source);
}
}
diff --git a/mlir/test/Dialect/Vector/vector-sink.mlir b/mlir/test/Dialect/Vector/vector-sink.mlir
index 6634f5d20cc85..59e181b7e1516 100644
--- a/mlir/test/Dialect/Vector/vector-sink.mlir
+++ b/mlir/test/Dialect/Vector/vector-sink.mlir
@@ -94,10 +94,10 @@ func.func @broadcast_vector_scalable(%arg1: vector<[4]xf32>, %arg2: vector<[4]xf
// CHECK-LABEL: func.func @broadcast_scalar_and_vec(
// CHECK-SAME: %[[ARG1:.*]]: index,
// CHECK-SAME: %[[ARG2:.*]]: vector<4xindex>) -> vector<1x4xindex> {
-// CHECK: %[[SPLAT:.*]] = vector.broadcast %[[ARG1]] : index to vector<1x4xindex>
-// CHECK: %[[BCAST:.*]] = vector.broadcast %[[ARG2]] : vector<4xindex> to vector<1x4xindex>
-// CHECK: %[[ADD:.*]] = arith.addi %[[SPLAT]], %[[BCAST]] : vector<1x4xindex>
-// CHECK: return %[[ADD]] : vector<1x4xindex>
+// CHECK: %[[SPLAT:.*]] = vector.broadcast %[[ARG1]] : index to vector<4xindex>
+// CHECK: %[[ADD:.*]] = arith.addi %[[SPLAT]], %[[ARG2]] : vector<4xindex>
+// CHECK: %[[BCAST:.*]] = vector.broadcast %[[ADD]] : vector<4xindex> to vector<1x4xindex>
+// CHECK: return %[[BCAST]] : vector<1x4xindex>
func.func @broadcast_scalar_and_vec(%arg1: index, %arg2: vector<4xindex>) -> vector<1x4xindex> {
%0 = vector.broadcast %arg1 : index to vector<1x4xindex>
%1 = vector.broadcast %arg2 : vector<4xindex> to vector<1x4xindex>
@@ -108,10 +108,10 @@ func.func @broadcast_scalar_and_vec(%arg1: index, %arg2: vector<4xindex>) -> vec
// CHECK-LABEL: func.func @broadcast_scalar_and_vec_scalable(
// CHECK-SAME: %[[ARG1:.*]]: index,
// CHECK-SAME: %[[ARG2:.*]]: vector<[4]xindex>) -> vector<1x[4]xindex> {
-// CHECK: %[[SPLAT:.*]] = vector.broadcast %[[ARG1]] : index to vector<1x[4]xindex>
-// CHECK: %[[BCAST:.*]] = vector.broadcast %[[ARG2]] : vector<[4]xindex> to vector<1x[4]xindex>
-// CHECK: %[[ADD:.*]] = arith.addi %[[SPLAT]], %[[BCAST]] : vector<1x[4]xindex>
-// CHECK: return %[[ADD]] : vector<1x[4]xindex>
+// CHECK: %[[SPLAT:.*]] = vector.broadcast %[[ARG1]] : index to vector<[4]xindex>
+// CHECK: %[[ADD:.*]] = arith.addi %[[SPLAT]], %[[ARG2]] : vector<[4]xindex>
+// CHECK: %[[BCAST:.*]] = vector.broadcast %[[ADD]] : vector<[4]xindex> to vector<1x[4]xindex>
+// CHECK: return %[[BCAST]] : vector<1x[4]xindex>
func.func @broadcast_scalar_and_vec_scalable(%arg1: index, %arg2: vector<[4]xindex>) -> vector<1x[4]xindex> {
%0 = vector.broadcast %arg1 : index to vector<1x[4]xindex>
%1 = vector.broadcast %arg2 : vector<[4]xindex> to vector<1x[4]xindex>
@@ -195,6 +195,21 @@ func.func @source_and_result_mismatch(%arg0 : f32) -> vector<1xi1> {
// -----
+// CHECK-LABEL: func.func @broadcast_vector_and_scalar_cmpf(
+// CHECK-SAME: %[[ARG0:.*]]: f32, %[[ARG1:.*]]: vector<4xf32>) -> vector<1x4xi1>
+// CHECK: %[[ARG0_BCAST:.*]] = vector.broadcast %[[ARG0]] : f32 to vector<4xf32>
+// CHECK: %[[CMP:.*]] = arith.cmpf uno, %[[ARG0_BCAST]], %[[ARG1]] : vector<4xf32>
+// CHECK: %[[BCAST:.*]] = vector.broadcast %[[CMP]] : vector<4xi1> to vector<1x4xi1>
+// CHECK: return %[[BCAST]] : vector<1x4xi1>
+func.func @broadcast_vector_and_scalar_cmpf(%arg0 : f32, %arg1 : vector<4xf32>) -> vector<1x4xi1> {
+ %0 = vector.broadcast %arg0 : f32 to vector<1x4xf32>
+ %1 = vector.broadcast %arg1 : vector<4xf32> to vector<1x4xf32>
+ %2 = arith.cmpf uno, %0, %1 : vector<1x4xf32>
+ return %2 : vector<1x4xi1>
+}
+
+// -----
+
// vector.fma only supports vectors - when the broadcast source is a scalar,
// currently it's not possible to replace this with e.g.:
// %scalar_res = vector.fma %scalar_1, %scalar2
@@ -214,14 +229,13 @@ func.func @negative_fma_scalar_broadcast_source(%arg0 : f32) -> vector<1xf32> {
// -----
-// CHECK-LABEL: func.func @negative_fma_mixed_scalar_and_vector_broadcast_source(
+// CHECK-LABEL: func.func @fma_mixed_scalar_and_vector_broadcast_source(
// CHECK-SAME: %[[ARG0:.*]]: vector<4xf32>, %[[ARG1:.*]]: f32, %[[ARG2:.*]]: vector<4xf32>)
-// CHECK: %[[BCAST0:.*]] = vector.broadcast %[[ARG0]] : vector<4xf32> to vector<3x4xf32>
-// CHECK: %[[BCAST1:.*]] = vector.broadcast %[[ARG1]] : f32 to vector<3x4xf32>
-// CHECK: %[[BCAST2:.*]] = vector.broadcast %[[ARG2]] : vector<4xf32> to vector<3x4xf32>
-// CHECK: %[[FMA:.*]] = vector.fma %[[BCAST0]], %[[BCAST1]], %[[BCAST2]] : vector<3x4xf32>
-// CHECK: return %[[FMA]] : vector<3x4xf32>
-func.func @negative_fma_mixed_scalar_and_vector_broadcast_source(%arg0: vector<4xf32>, %arg1: f32, %arg2: vector<4xf32>) -> vector<3x4xf32> {
+// CHECK: %[[ARG1_BCAST:.*]] = vector.broadcast %[[ARG1]] : f32 to vector<4xf32>
+// CHECK: %[[FMA:.*]] = vector.fma %[[ARG0]], %[[ARG1_BCAST]], %[[ARG2]] : vector<4xf32>
+// CHECK: %[[BCAST:.*]] = vector.broadcast %[[FMA]] : vector<4xf32> to vector<3x4xf32>
+// CHECK: return %[[BCAST]] : vector<3x4xf32>
+func.func @fma_mixed_scalar_and_vector_broadcast_source(%arg0: vector<4xf32>, %arg1: f32, %arg2: vector<4xf32>) -> vector<3x4xf32> {
%0 = vector.broadcast %arg0 : vector<4xf32> to vector<3x4xf32>
%1 = vector.broadcast %arg1 : f32 to vector<3x4xf32>
%2 = vector.broadcast %arg2 : vector<4xf32> to vector<3x4xf32>
@@ -231,8 +245,27 @@ func.func @negative_fma_mixed_scalar_and_vector_broadcast_source(%arg0: vector<4
// -----
+// The vector source may occur after a scalar source. It still determines the
+// shape used for the smaller FMA.
+
+// CHECK-LABEL: func.func @fma_mixed_scalar_first_broadcast_source(
+// CHECK-SAME: %[[ARG0:.*]]: f32, %[[ARG1:.*]]: vector<4xf32>, %[[ARG2:.*]]: vector<4xf32>)
+// CHECK: %[[ARG0_BCAST:.*]] = vector.broadcast %[[ARG0]] : f32 to vector<4xf32>
+// CHECK: %[[FMA:.*]] = vector.fma %[[ARG0_BCAST]], %[[ARG1]], %[[ARG2]] : vector<4xf32>
+// CHECK: %[[BCAST:.*]] = vector.broadcast %[[FMA]] : vector<4xf32> to vector<3x4xf32>
+// CHECK: return %[[BCAST]] : vector<3x4xf32>
+func.func @fma_mixed_scalar_first_broadcast_source(%arg0: f32, %arg1: vector<4xf32>, %arg2: vector<4xf32>) -> vector<3x4xf32> {
+ %0 = vector.broadcast %arg0 : f32 to vector<3x4xf32>
+ %1 = vector.broadcast %arg1 : vector<4xf32> to vector<3x4xf32>
+ %2 = vector.broadcast %arg2 : vector<4xf32> to vector<3x4xf32>
+ %3 = vector.fma %0, %1, %2 : vector<3x4xf32>
+ return %3 : vector<3x4xf32>
+}
+
+// -----
+
// vector.fma only supports vector operands, hence the broadcast can only be
-// sunk when the broadcast source is a vector as well.
+// sunk when there is a vector source to determine the smaller vector type.
// CHECK-LABEL: func.func @fma_vector_broadcast_source(
// CHECK-SAME: %[[ARG0:.*]]: vector<4xf32>, %[[ARG1:.*]]: vector<4xf32>, %[[ARG2:.*]]: vector<4xf32>)
More information about the Mlir-commits
mailing list