[Mlir-commits] [mlir] [mlir] [vector] Reorder element/broadcast for FMA if the source is not scalar (PR #211208)
Chuanqi Xu
llvmlistbot at llvm.org
Wed Jul 22 02:16:49 PDT 2026
https://github.com/ChuanqiXu9 created https://github.com/llvm/llvm-project/pull/211208
The diagnostic says the transform is not performed as the source may be a scalar. But we can make it if we can check the source is not a scalar.
AI assisted.
>From ca4a8c504c93d837cfa053a1b86b1cbd0167a728 Mon Sep 17 00:00:00 2001
From: "yedeng.yd" <yedeng.yd at alibaba-inc.com>
Date: Wed, 22 Jul 2026 16:56:46 +0800
Subject: [PATCH] [mlir] [vector] Reorder element/broadcast for FMA if the
source is not scalar
The diagnostic says the transform is not performed as the source may be
a scalar. But we can make it if we can check the source is not a scalar.
AI assisted.
---
.../Vector/Transforms/VectorTransforms.cpp | 17 ++--
mlir/test/Dialect/Vector/vector-sink.mlir | 81 +++++++++++++++++--
2 files changed, 85 insertions(+), 13 deletions(-)
diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
index 17e6ad640cada..e1d9bc8d0f4a0 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
@@ -1043,12 +1043,6 @@ struct ReorderElementwiseOpsOnBroadcast final
op, "Op doesn't have ElementwiseMappableTraits");
if (op->getNumOperands() == 0)
return failure();
- if (isa<vector::FMAOp>(op)) {
- return rewriter.notifyMatchFailure(
- op,
- "Op only accepts vector types - not supported as broadcast source "
- "might be a scalar");
- }
Type resultElemType = resultType.getElementType();
@@ -1068,6 +1062,17 @@ struct ReorderElementwiseOpsOnBroadcast final
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: Support scalar sources for `vector.fma` by rewriting to
+ // `math.fma`.
+ 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`).
diff --git a/mlir/test/Dialect/Vector/vector-sink.mlir b/mlir/test/Dialect/Vector/vector-sink.mlir
index 69fba88a14048..b52828419b10a 100644
--- a/mlir/test/Dialect/Vector/vector-sink.mlir
+++ b/mlir/test/Dialect/Vector/vector-sink.mlir
@@ -195,17 +195,18 @@ func.func @source_and_result_mismatch(%arg0 : f32) -> vector<1xi1> {
// -----
-// vector.fma only supports vectors - currently it's not possible to replace this with e.g.:
+// 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
// %vec_res = vector.broadcast %scalar_res
//
-// TODO: It should be possible to support this case
+// TODO: It should be possible to support this case by rewriting to math.fma
-// CHECK-LABEL: func.func @negative_op_only_supports_vectors
- // CHECK: %[[BROADCAST:.+]] = vector.broadcast
- // CHECK: %[[RESULT:.+]] = vector.fma %[[BROADCAST]]
- // CHECK: return %[[RESULT]]
-func.func @negative_op_only_supports_vectors(%arg0 : f32) -> vector<1xf32> {
+// CHECK-LABEL: func.func @negative_fma_scalar_broadcast_source
+// CHECK: %[[BROADCAST:.+]] = vector.broadcast
+// CHECK: %[[RESULT:.+]] = vector.fma %[[BROADCAST]]
+// CHECK: return %[[RESULT]]
+func.func @negative_fma_scalar_broadcast_source(%arg0 : f32) -> vector<1xf32> {
%0 = vector.broadcast %arg0 : f32 to vector<1xf32>
%1 = vector.fma %0, %0, %0 : vector<1xf32>
return %1 : vector<1xf32>
@@ -213,6 +214,72 @@ func.func @negative_op_only_supports_vectors(%arg0 : f32) -> vector<1xf32> {
// -----
+// CHECK-LABEL: func.func @negative_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> {
+ %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>
+ 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.
+
+// CHECK-LABEL: func.func @fma_vector_broadcast_source(
+// CHECK-SAME: %[[ARG0:.*]]: vector<4xf32>, %[[ARG1:.*]]: vector<4xf32>, %[[ARG2:.*]]: vector<4xf32>)
+// CHECK: %[[FMA:.*]] = vector.fma %[[ARG0]], %[[ARG1]], %[[ARG2]] : vector<4xf32>
+// CHECK: %[[BCAST:.*]] = vector.broadcast %[[FMA]] : vector<4xf32> to vector<3x4xf32>
+// CHECK: return %[[BCAST]] : vector<3x4xf32>
+func.func @fma_vector_broadcast_source(%arg0: vector<4xf32>, %arg1: vector<4xf32>, %arg2: vector<4xf32>) -> vector<3x4xf32> {
+ %0 = vector.broadcast %arg0 : vector<4xf32> 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>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @fma_vector_broadcast_source_scalable(
+// CHECK-SAME: %[[ARG0:.*]]: vector<[4]xf32>, %[[ARG1:.*]]: vector<[4]xf32>, %[[ARG2:.*]]: vector<[4]xf32>)
+// CHECK: %[[FMA:.*]] = vector.fma %[[ARG0]], %[[ARG1]], %[[ARG2]] : vector<[4]xf32>
+// CHECK: %[[BCAST:.*]] = vector.broadcast %[[FMA]] : vector<[4]xf32> to vector<3x[4]xf32>
+// CHECK: return %[[BCAST]] : vector<3x[4]xf32>
+func.func @fma_vector_broadcast_source_scalable(%arg0: vector<[4]xf32>, %arg1: vector<[4]xf32>, %arg2: vector<[4]xf32>) -> vector<3x[4]xf32> {
+ %0 = vector.broadcast %arg0 : vector<[4]xf32> to vector<3x[4]xf32>
+ %1 = vector.broadcast %arg1 : vector<[4]xf32> to vector<3x[4]xf32>
+ %2 = vector.broadcast %arg2 : vector<[4]xf32> to vector<3x[4]xf32>
+ %3 = vector.fma %0, %1, %2 : vector<3x[4]xf32>
+ return %3 : vector<3x[4]xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @fma_vector_broadcast_source_and_splat_const(
+// CHECK-SAME: %[[ARG0:.*]]: vector<4xf32>, %[[ARG1:.*]]: vector<4xf32>)
+// CHECK: %[[NEW_CST:.*]] = arith.constant dense<2.000000e+00> : vector<4xf32>
+// CHECK: %[[FMA:.*]] = vector.fma %[[ARG0]], %[[ARG1]], %[[NEW_CST]] : vector<4xf32>
+// CHECK: %[[BCAST:.*]] = vector.broadcast %[[FMA]] : vector<4xf32> to vector<3x4xf32>
+// CHECK: return %[[BCAST]] : vector<3x4xf32>
+func.func @fma_vector_broadcast_source_and_splat_const(%arg0: vector<4xf32>, %arg1: vector<4xf32>) -> vector<3x4xf32> {
+ %0 = vector.broadcast %arg0 : vector<4xf32> to vector<3x4xf32>
+ %1 = vector.broadcast %arg1 : vector<4xf32> to vector<3x4xf32>
+ %cst = arith.constant dense<2.0> : vector<3x4xf32>
+ %2 = vector.fma %0, %1, %cst : vector<3x4xf32>
+ return %2 : vector<3x4xf32>
+}
+
+// -----
+
// CHECK-LABEL: func.func @broadcast_scalar_and_splat_const(
// CHECK-SAME: %[[ARG_0:.*]]: index) -> vector<1x4xindex> {
// CHECK: %[[NEW_CST:.*]] = arith.constant 2 : index
More information about the Mlir-commits
mailing list