[Mlir-commits] [mlir] 3de7874 - [mlir] [vector] try promoting scalar when reordering broadcast/elementwise (#212180)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Aug 24 19:55:20 PDT 2026
Author: Chuanqi Xu
Date: 2026-08-25T10:55:14+08:00
New Revision: 3de7874ccbcce9b5eb53dab392c4b889c4fb3567
URL: https://github.com/llvm/llvm-project/commit/3de7874ccbcce9b5eb53dab392c4b889c4fb3567
DIFF: https://github.com/llvm/llvm-project/commit/3de7874ccbcce9b5eb53dab392c4b889c4fb3567.diff
LOG: [mlir] [vector] try promoting scalar when reordering broadcast/elementwise (#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.
Added:
Modified:
mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
mlir/test/Dialect/Linalg/vectorization/extract-with-patterns.mlir
mlir/test/Dialect/Vector/vector-sink.mlir
Removed:
################################################################################
diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
index d7beb08d3333c..eeb2c7259b6e2 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
@@ -1046,41 +1046,56 @@ struct ReorderElementwiseOpsOnBroadcast final
Type resultElemType = resultType.getElementType();
- // Get the type of the first non-constant operand
+ // 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()) {
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 all non-constant operands are scalar, choose the first source.
+ 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 +1123,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/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]]
diff --git a/mlir/test/Dialect/Vector/vector-sink.mlir b/mlir/test/Dialect/Vector/vector-sink.mlir
index 6634f5d20cc85..ac1e4fa1ea1cf 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,24 @@ 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>
+// 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 +232,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 +248,30 @@ func.func @negative_fma_mixed_scalar_and_vector_broadcast_source(%arg0: vector<4
// -----
+// 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>)
+// 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