[Mlir-commits] [mlir] [mlir] [vector] try promoting scalar when reordering broadcast/elementwise (PR #212180)
Chuanqi Xu
llvmlistbot at llvm.org
Mon Aug 24 01:49:17 PDT 2026
https://github.com/ChuanqiXu9 updated https://github.com/llvm/llvm-project/pull/212180
>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 1/3] [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>)
>From b2b8504010bfa4f6ee9027c88ea788dc37f8e450 Mon Sep 17 00:00:00 2001
From: "yedeng.yd" <yedeng.yd at alibaba-inc.com>
Date: Mon, 27 Jul 2026 16:05:25 +0800
Subject: [PATCH 2/3] update test
---
.../Dialect/Linalg/vectorization/extract-with-patterns.mlir | 6 +++---
1 file changed, 3 insertions(+), 3 deletions(-)
diff --git a/mlir/test/Dialect/Linalg/vectorization/extract-with-patterns.mlir b/mlir/test/Dialect/Linalg/vectorization/extract-with-patterns.mlir
index e04a3f1a83d35..45f8f6a393f0a 100644
--- a/mlir/test/Dialect/Linalg/vectorization/extract-with-patterns.mlir
+++ b/mlir/test/Dialect/Linalg/vectorization/extract-with-patterns.mlir
@@ -516,9 +516,9 @@ func.func @vectorize_reverse_like_tensor_extract(%arg0: tensor<1x2x3xf32>, %arg1
// CHECK-DAG: %[[PASSTHRU:.*]] = arith.constant dense<0.000000e+00> : vector<1x1x3xf32>
// CHECK-DAG: %[[INIT_IDX:.+]] = arith.constant dense<[2, 1, 0]> : vector<3xindex>
// CHECK: %[[T0:.+]] = arith.muli %[[ARG2]], %[[C3]] : index
-// CHECK: %[[T1:.+]] = vector.broadcast %[[T0]] : index to vector<1x1x3xindex>
-// CHECK: %[[T2:.+]] = vector.broadcast %[[INIT_IDX]]
-// CHECK: %[[T3:.+]] = arith.addi %[[T2]], %[[T1]]
+// CHECK: %[[T1:.+]] = vector.broadcast %[[T0]] : index to vector<3xindex>
+// CHECK: %[[T2:.+]] = arith.addi %[[T1]], %[[INIT_IDX]] : vector<3xindex>
+// CHECK: %[[T3:.+]] = vector.broadcast %[[T2]] : vector<3xindex> to vector<1x1x3xindex>
// CHECK: %[[GATHER:.*]] = vector.gather %[[ARG0]][%[[C0]], %[[C0]], %[[C0]]] [%[[T3]]], %[[MASK]], %[[PASSTHRU]]
// CHECK: vector.transfer_write %[[GATHER]]
>From 05182a6a2786fae2d5d6c9e70e7598f6f011365e Mon Sep 17 00:00:00 2001
From: Chuanqi Xu <chuanqi.xcq at alibaba-inc.com>
Date: Mon, 24 Aug 2026 16:48:42 +0800
Subject: [PATCH 3/3] Update
---
.../lib/Dialect/Vector/Transforms/VectorTransforms.cpp | 8 +++++---
mlir/test/Dialect/Vector/vector-sink.mlir | 10 ++++++++--
2 files changed, 13 insertions(+), 5 deletions(-)
diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
index b2626d17b8bbd..eeb2c7259b6e2 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
@@ -1046,9 +1046,10 @@ struct ReorderElementwiseOpsOnBroadcast final
Type resultElemType = resultType.getElementType();
- // 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.
+ // Select the source shape for the reordered computation. Prefer the first
+ // non-constant vector source so that scalar sources can be broadcast to its
+ // shape. The compatibility check below ensures that all vector sources have
+ // the same shape and scalable dimensions.
Value broadcastSource;
Value firstBroadcastSource;
for (Value operand : op->getOperands()) {
@@ -1067,6 +1068,7 @@ struct ReorderElementwiseOpsOnBroadcast final
break;
}
}
+ // If all non-constant operands are scalar, choose the first source.
if (!broadcastSource)
broadcastSource = firstBroadcastSource;
if (!broadcastSource)
diff --git a/mlir/test/Dialect/Vector/vector-sink.mlir b/mlir/test/Dialect/Vector/vector-sink.mlir
index 59e181b7e1516..ac1e4fa1ea1cf 100644
--- a/mlir/test/Dialect/Vector/vector-sink.mlir
+++ b/mlir/test/Dialect/Vector/vector-sink.mlir
@@ -195,6 +195,9 @@ func.func @source_and_result_mismatch(%arg0 : f32) -> vector<1xi1> {
// -----
+// Verify that the elementwise computation is performed on the lowest-rank
+// vector and only the result is broadcast to the original shape.
+//
// 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>
@@ -245,8 +248,11 @@ func.func @fma_mixed_scalar_and_vector_broadcast_source(%arg0: vector<4xf32>, %a
// -----
-// The vector source may occur after a scalar source. It still determines the
-// shape used for the smaller FMA.
+// Verify that the elementwise computation is performed on the lowest-rank
+// vector and only the result is broadcast to the original shape.
+//
+// The vector source may occur after a scalar source. The later vector source
+// 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>)
More information about the Mlir-commits
mailing list