[Mlir-commits] [mlir] [mlir][vector] Clarify the semantics of BroadcastOp (PR #101928)
Andrzej Warzyński
llvmlistbot at llvm.org
Tue Aug 6 02:45:04 PDT 2024
https://github.com/banach-space updated https://github.com/llvm/llvm-project/pull/101928
>From 7c08a8bfd434d0ce1b71e891cc027fc98542353c Mon Sep 17 00:00:00 2001
From: Andrzej Warzynski <andrzej.warzynski at arm.com>
Date: Sun, 4 Aug 2024 15:31:49 +0100
Subject: [PATCH 1/3] [mlir][vector] Clarify the semantics of BroadcastOp
Clarifies the semantics of `vector.broadcast` in the context of scalable
vectors. In particular, broadcasting a unit scalable dim, `[1]`, is not
valid unless there's a match between the output and the input dims.
See the examples below for an illustration:
```mlir
// VALID
%0 = vector.broadcast %arg0 : vector<[1]xf32> to vector<4x[1]xf32>
// INVALID
%0 = vector.broadcast %arg0 : vector<[1]xf32> to vector<[4]xf32>
// VALID FIXED-WIDTH EQUIVALENT
%0 = vector.broadcast %arg0 : vector<1xf32> to vector<4xf32>
```
Documentation, the Op verifier and tests are updated accordingly.
---
.../mlir/Dialect/Vector/IR/VectorOps.h | 6 ++-
.../mlir/Dialect/Vector/IR/VectorOps.td | 2 +
mlir/lib/Dialect/Vector/IR/VectorOps.cpp | 45 ++++++++++++++-----
mlir/test/Dialect/Vector/invalid.mlir | 14 ++++++
4 files changed, 56 insertions(+), 11 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.h b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.h
index ac55433fadb2f..9f61f7c866d3d 100644
--- a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.h
+++ b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.h
@@ -68,9 +68,13 @@ enum class BroadcastableToResult {
DimensionMismatch = 2,
SourceTypeNotAVector = 3
};
+struct VectorDim {
+ int64_t dim;
+ bool scalableFlag;
+};
BroadcastableToResult
isBroadcastableTo(Type srcType, VectorType dstVectorType,
- std::pair<int, int> *mismatchingDims = nullptr);
+ std::pair<VectorDim, VectorDim> *mismatchingDims = nullptr);
/// Collect a set of vector-to-vector canonicalization patterns.
void populateVectorToVectorCanonicalizationPatterns(RewritePatternSet &patterns,
diff --git a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
index 434ff3956c250..08bff3d5e1382 100644
--- a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
+++ b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
@@ -367,6 +367,8 @@ def Vector_BroadcastOp :
s_1 x .. x s_j x .. x s_k
<duplication> <potential stretch>
```
+ * a scalable unit dimeension, `[1]`, must match exactly.
+
The source operand is duplicated over all the missing leading dimensions
and stretched over the trailing dimensions where the source has a non-equal
dimension of 1. These rules imply that any scalar broadcast (k=0) to any
diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index 5047bd925d4c5..673c128932893 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -2371,9 +2371,9 @@ Value BroadcastOp::createOrFoldBroadcastOp(
return res;
}
-BroadcastableToResult
-mlir::vector::isBroadcastableTo(Type srcType, VectorType dstVectorType,
- std::pair<int, int> *mismatchingDims) {
+BroadcastableToResult mlir::vector::isBroadcastableTo(
+ Type srcType, VectorType dstVectorType,
+ std::pair<VectorDim, VectorDim> *mismatchingDims) {
// Broadcast scalar to vector of the same element type.
if (srcType.isIntOrIndexOrFloat() && dstVectorType &&
getElementTypeOrSelf(srcType) == getElementTypeOrSelf(dstVectorType))
@@ -2391,12 +2391,28 @@ mlir::vector::isBroadcastableTo(Type srcType, VectorType dstVectorType,
// (all leading dimensions are simply duplicated).
int64_t lead = dstRank - srcRank;
for (int64_t r = 0; r < srcRank; ++r) {
+ bool mismatch = false;
+
+ // Check fixed-width dims
int64_t srcDim = srcVectorType.getDimSize(r);
int64_t dstDim = dstVectorType.getDimSize(lead + r);
- if (srcDim != 1 && srcDim != dstDim) {
+ if ((srcDim != 1 && srcDim != dstDim))
+ mismatch = true;
+
+ // Check scalable flags
+ bool srcDimScalableFlag = srcVectorType.getScalableDims()[r];
+ bool dstDimScalableFlag = dstVectorType.getScalableDims()[lead + r];
+ if ((srcDim == 1 && srcDimScalableFlag && dstDim != 1) ||
+ (srcDimScalableFlag && !dstDimScalableFlag))
+ mismatch = true;
+
+ if (mismatch) {
if (mismatchingDims) {
- mismatchingDims->first = srcDim;
- mismatchingDims->second = dstDim;
+ mismatchingDims->first.dim = srcDim;
+ mismatchingDims->first.scalableFlag = srcDimScalableFlag;
+
+ mismatchingDims->second.dim = dstDim;
+ mismatchingDims->second.scalableFlag = dstDimScalableFlag;
}
return BroadcastableToResult::DimensionMismatch;
}
@@ -2406,16 +2422,25 @@ mlir::vector::isBroadcastableTo(Type srcType, VectorType dstVectorType,
}
LogicalResult BroadcastOp::verify() {
- std::pair<int, int> mismatchingDims;
+ std::pair<VectorDim, VectorDim> mismatchingDims;
BroadcastableToResult res = isBroadcastableTo(
getSourceType(), getResultVectorType(), &mismatchingDims);
if (res == BroadcastableToResult::Success)
return success();
if (res == BroadcastableToResult::SourceRankHigher)
return emitOpError("source rank higher than destination rank");
- if (res == BroadcastableToResult::DimensionMismatch)
- return emitOpError("dimension mismatch (")
- << mismatchingDims.first << " vs. " << mismatchingDims.second << ")";
+ if (res == BroadcastableToResult::DimensionMismatch) {
+ std::string msg =
+ (Twine("dimension mismatch (") +
+ (mismatchingDims.first.scalableFlag ? "[" : "") +
+ std::to_string(mismatchingDims.first.dim) +
+ (mismatchingDims.first.scalableFlag ? "]" : "") + " vs. " +
+ (mismatchingDims.second.scalableFlag ? "[" : "") +
+ std::to_string(mismatchingDims.second.dim) +
+ (mismatchingDims.second.scalableFlag ? "]" : "") + ")")
+ .str();
+ return emitOpError(msg);
+ }
if (res == BroadcastableToResult::SourceTypeNotAVector)
return emitOpError("source type is not a vector");
llvm_unreachable("unexpected vector.broadcast op error");
diff --git a/mlir/test/Dialect/Vector/invalid.mlir b/mlir/test/Dialect/Vector/invalid.mlir
index 00914c1d1baf6..6dd690be032c7 100644
--- a/mlir/test/Dialect/Vector/invalid.mlir
+++ b/mlir/test/Dialect/Vector/invalid.mlir
@@ -35,6 +35,20 @@ func.func @broadcast_dim2_mismatch(%arg0: vector<4x8xf32>) {
// -----
+func.func @broadcast_scalable_unit_dim(%arg0: vector<[1]xf32>) {
+ // expected-error at +1 {{'vector.broadcast' op dimension mismatch ([1] vs. [4])}}
+ %0 = vector.broadcast %arg0 : vector<[1]xf32> to vector<[4]xf32>
+}
+
+// -----
+
+func.func @broadcast_scalable_to_fixed(%arg0: vector<[1]xf32>) {
+ // expected-error at +1 {{'vector.broadcast' op dimension mismatch ([1] vs. 1)}}
+ %0 = vector.broadcast %arg0 : vector<[1]xf32> to vector<4x1xf32>
+}
+
+// -----
+
func.func @broadcast_unknown(%arg0: memref<4x8xf32>) {
// expected-error at +1 {{'vector.broadcast' op source type is not a vector}}
%1 = vector.broadcast %arg0 : memref<4x8xf32> to vector<1x8xf32>
>From 74d843c3635d3192060294f77ba65bf6151bbfa9 Mon Sep 17 00:00:00 2001
From: Andrzej Warzynski <andrzej.warzynski at arm.com>
Date: Mon, 5 Aug 2024 14:46:02 +0100
Subject: [PATCH 2/3] fixup! [mlir][vector] Clarify the semantics of
BroadcastOp
Address comments from Hugo - thank you!
---
mlir/include/mlir/Dialect/Vector/IR/VectorOps.td | 2 +-
mlir/lib/Dialect/Vector/IR/VectorOps.cpp | 2 +-
mlir/test/Dialect/Vector/invalid.mlir | 7 +++++++
3 files changed, 9 insertions(+), 2 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
index 08bff3d5e1382..456ca25f96cc2 100644
--- a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
+++ b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
@@ -367,7 +367,7 @@ def Vector_BroadcastOp :
s_1 x .. x s_j x .. x s_k
<duplication> <potential stretch>
```
- * a scalable unit dimeension, `[1]`, must match exactly.
+ * a scalable unit dimension, `[1]`, must match exactly.
The source operand is duplicated over all the missing leading dimensions
and stretched over the trailing dimensions where the source has a non-equal
diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index 673c128932893..49b891d068457 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -2403,7 +2403,7 @@ BroadcastableToResult mlir::vector::isBroadcastableTo(
bool srcDimScalableFlag = srcVectorType.getScalableDims()[r];
bool dstDimScalableFlag = dstVectorType.getScalableDims()[lead + r];
if ((srcDim == 1 && srcDimScalableFlag && dstDim != 1) ||
- (srcDimScalableFlag && !dstDimScalableFlag))
+ (srcDimScalableFlag != dstDimScalableFlag))
mismatch = true;
if (mismatch) {
diff --git a/mlir/test/Dialect/Vector/invalid.mlir b/mlir/test/Dialect/Vector/invalid.mlir
index 6dd690be032c7..13712578536cd 100644
--- a/mlir/test/Dialect/Vector/invalid.mlir
+++ b/mlir/test/Dialect/Vector/invalid.mlir
@@ -42,6 +42,13 @@ func.func @broadcast_scalable_unit_dim(%arg0: vector<[1]xf32>) {
// -----
+func.func @broadcast_fixed_to_scalable(%arg0: vector<2xf32>) {
+ // expected-error at +1 {{'vector.broadcast' op dimension mismatch (2 vs. [2])}}
+ %0 = vector.broadcast %arg0 : vector<2xf32> to vector<[2]xf32>
+}
+
+// -----
+
func.func @broadcast_scalable_to_fixed(%arg0: vector<[1]xf32>) {
// expected-error at +1 {{'vector.broadcast' op dimension mismatch ([1] vs. 1)}}
%0 = vector.broadcast %arg0 : vector<[1]xf32> to vector<4x1xf32>
>From 587f08e445f68a8311c2c5bdaa967d80fea34211 Mon Sep 17 00:00:00 2001
From: Andrzej Warzynski <andrzej.warzynski at arm.com>
Date: Tue, 6 Aug 2024 10:44:34 +0100
Subject: [PATCH 3/3] fixup! fixup! [mlir][vector] Clarify the semantics of
BroadcastOp
Address comments from Jakub
---
.../mlir/Dialect/Vector/IR/VectorOps.h | 3 +-
.../mlir/Dialect/Vector/IR/VectorOps.td | 2 +-
mlir/lib/Dialect/Vector/IR/VectorOps.cpp | 30 +++++++++----------
3 files changed, 18 insertions(+), 17 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.h b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.h
index 9f61f7c866d3d..64bfb0fab743a 100644
--- a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.h
+++ b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.h
@@ -68,9 +68,10 @@ enum class BroadcastableToResult {
DimensionMismatch = 2,
SourceTypeNotAVector = 3
};
+
struct VectorDim {
int64_t dim;
- bool scalableFlag;
+ bool isScalable;
};
BroadcastableToResult
isBroadcastableTo(Type srcType, VectorType dstVectorType,
diff --git a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
index 456ca25f96cc2..ead7faae46cbe 100644
--- a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
+++ b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
@@ -367,7 +367,7 @@ def Vector_BroadcastOp :
s_1 x .. x s_j x .. x s_k
<duplication> <potential stretch>
```
- * a scalable unit dimension, `[1]`, must match exactly.
+ * in addition, any scalable unit dimension, `[1]`, must match exactly.
The source operand is duplicated over all the missing leading dimensions
and stretched over the trailing dimensions where the source has a non-equal
diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index 49b891d068457..b8e2fc0b6449e 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -2390,29 +2390,29 @@ BroadcastableToResult mlir::vector::isBroadcastableTo(
// Source has an exact match or singleton value for all trailing dimensions
// (all leading dimensions are simply duplicated).
int64_t lead = dstRank - srcRank;
- for (int64_t r = 0; r < srcRank; ++r) {
+ for (int64_t dimIdx = 0; dimIdx < srcRank; ++dimIdx) {
bool mismatch = false;
- // Check fixed-width dims
- int64_t srcDim = srcVectorType.getDimSize(r);
- int64_t dstDim = dstVectorType.getDimSize(lead + r);
- if ((srcDim != 1 && srcDim != dstDim))
+ // Check fixed-width dims.
+ int64_t srcDim = srcVectorType.getDimSize(dimIdx);
+ int64_t dstDim = dstVectorType.getDimSize(lead + dimIdx);
+ if (srcDim != 1 && srcDim != dstDim)
mismatch = true;
- // Check scalable flags
- bool srcDimScalableFlag = srcVectorType.getScalableDims()[r];
- bool dstDimScalableFlag = dstVectorType.getScalableDims()[lead + r];
+ // Check scalable flags.
+ bool srcDimScalableFlag = srcVectorType.getScalableDims()[dimIdx];
+ bool dstDimScalableFlag = dstVectorType.getScalableDims()[lead + dimIdx];
if ((srcDim == 1 && srcDimScalableFlag && dstDim != 1) ||
(srcDimScalableFlag != dstDimScalableFlag))
mismatch = true;
if (mismatch) {
- if (mismatchingDims) {
+ if (mismatchingDims != nullptr) {
mismatchingDims->first.dim = srcDim;
- mismatchingDims->first.scalableFlag = srcDimScalableFlag;
+ mismatchingDims->first.isScalable = srcDimScalableFlag;
mismatchingDims->second.dim = dstDim;
- mismatchingDims->second.scalableFlag = dstDimScalableFlag;
+ mismatchingDims->second.isScalable = dstDimScalableFlag;
}
return BroadcastableToResult::DimensionMismatch;
}
@@ -2432,12 +2432,12 @@ LogicalResult BroadcastOp::verify() {
if (res == BroadcastableToResult::DimensionMismatch) {
std::string msg =
(Twine("dimension mismatch (") +
- (mismatchingDims.first.scalableFlag ? "[" : "") +
+ (mismatchingDims.first.isScalable ? "[" : "") +
std::to_string(mismatchingDims.first.dim) +
- (mismatchingDims.first.scalableFlag ? "]" : "") + " vs. " +
- (mismatchingDims.second.scalableFlag ? "[" : "") +
+ (mismatchingDims.first.isScalable ? "]" : "") + " vs. " +
+ (mismatchingDims.second.isScalable ? "[" : "") +
std::to_string(mismatchingDims.second.dim) +
- (mismatchingDims.second.scalableFlag ? "]" : "") + ")")
+ (mismatchingDims.second.isScalable ? "]" : "") + ")")
.str();
return emitOpError(msg);
}
More information about the Mlir-commits
mailing list