From 7c08a8bfd434d0ce1b71e891cc027fc98542353c Mon Sep 17 00:00:00 2001 From: Andrzej Warzynski Date: Sun, 4 Aug 2024 15:31:49 +0100 Subject: [PATCH 1/5] [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 *mismatchingDims = nullptr); + std::pair *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 ``` + * 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 *mismatchingDims) { +BroadcastableToResult mlir::vector::isBroadcastableTo( + Type srcType, VectorType dstVectorType, + std::pair *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 mismatchingDims; + std::pair 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@+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@+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@+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 Date: Mon, 5 Aug 2024 14:46:02 +0100 Subject: [PATCH 2/5] 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 ``` - * 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@+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@+1 {{'vector.broadcast' op dimension mismatch ([1] vs. 1)}} %0 = vector.broadcast %arg0 : vector<[1]xf32> to vector<4x1xf32> From 89be55b0bf8f996514dc4971f3d97ee246c77fbf Mon Sep 17 00:00:00 2001 From: Andrzej Warzynski Date: Tue, 6 Aug 2024 10:44:34 +0100 Subject: [PATCH 3/5] 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 | 39 +++++++++---------- 3 files changed, 22 insertions(+), 22 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 ``` - * 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..e9e708fbba500 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; } @@ -2430,15 +2430,14 @@ LogicalResult BroadcastOp::verify() { if (res == BroadcastableToResult::SourceRankHigher) return emitOpError("source rank higher than destination rank"); 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(); + std::string msg = (Twine("dimension mismatch (") + + (mismatchingDims.first.isScalable ? "[" : "") + + std::to_string(mismatchingDims.first.dim) + + (mismatchingDims.first.isScalable ? "]" : "") + " vs. " + + (mismatchingDims.second.isScalable ? "[" : "") + + std::to_string(mismatchingDims.second.dim) + + (mismatchingDims.second.isScalable ? "]" : "") + ")") + .str(); return emitOpError(msg); } if (res == BroadcastableToResult::SourceTypeNotAVector) From 001a3480a237c801c8025d9995358dcbd7229a81 Mon Sep 17 00:00:00 2001 From: Andrzej Warzynski Date: Tue, 6 Aug 2024 16:27:38 +0100 Subject: [PATCH 4/5] fixup! fixup! fixup! [mlir][vector] Clarify the semantics of BroadcastOp More comments and simplifications --- mlir/lib/Dialect/Vector/IR/VectorOps.cpp | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp index e9e708fbba500..2146678ec7995 100644 --- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp +++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp @@ -2391,22 +2391,24 @@ BroadcastableToResult mlir::vector::isBroadcastableTo( // (all leading dimensions are simply duplicated). int64_t lead = dstRank - srcRank; for (int64_t dimIdx = 0; dimIdx < srcRank; ++dimIdx) { - bool mismatch = false; + // Have mismatching dims (in the sense of vector.broadcast semantics) been + // encountered? + bool foundMismatchingDims = false; // Check fixed-width dims. int64_t srcDim = srcVectorType.getDimSize(dimIdx); int64_t dstDim = dstVectorType.getDimSize(lead + dimIdx); if (srcDim != 1 && srcDim != dstDim) - mismatch = true; + foundMismatchingDims = true; // Check scalable flags. bool srcDimScalableFlag = srcVectorType.getScalableDims()[dimIdx]; bool dstDimScalableFlag = dstVectorType.getScalableDims()[lead + dimIdx]; if ((srcDim == 1 && srcDimScalableFlag && dstDim != 1) || (srcDimScalableFlag != dstDimScalableFlag)) - mismatch = true; + foundMismatchingDims = true; - if (mismatch) { + if (foundMismatchingDims) { if (mismatchingDims != nullptr) { mismatchingDims->first.dim = srcDim; mismatchingDims->first.isScalable = srcDimScalableFlag; @@ -2432,10 +2434,10 @@ LogicalResult BroadcastOp::verify() { if (res == BroadcastableToResult::DimensionMismatch) { std::string msg = (Twine("dimension mismatch (") + (mismatchingDims.first.isScalable ? "[" : "") + - std::to_string(mismatchingDims.first.dim) + + Twine(mismatchingDims.first.dim) + (mismatchingDims.first.isScalable ? "]" : "") + " vs. " + (mismatchingDims.second.isScalable ? "[" : "") + - std::to_string(mismatchingDims.second.dim) + + Twine(mismatchingDims.second.dim) + (mismatchingDims.second.isScalable ? "]" : "") + ")") .str(); return emitOpError(msg); From ab4083183c37a92b7ee74023c5fdfb310bcae0e0 Mon Sep 17 00:00:00 2001 From: Andrzej Warzynski Date: Tue, 6 Aug 2024 17:32:07 +0100 Subject: [PATCH 5/5] fixup! fixup! fixup! fixup! [mlir][vector] Clarify the semantics of BroadcastOp Avoid using llvm::Twine --- mlir/lib/Dialect/Vector/IR/VectorOps.cpp | 16 +++++++--------- 1 file changed, 7 insertions(+), 9 deletions(-) diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp index 2146678ec7995..22220a6672382 100644 --- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp +++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp @@ -2432,15 +2432,13 @@ LogicalResult BroadcastOp::verify() { if (res == BroadcastableToResult::SourceRankHigher) return emitOpError("source rank higher than destination rank"); if (res == BroadcastableToResult::DimensionMismatch) { - std::string msg = (Twine("dimension mismatch (") + - (mismatchingDims.first.isScalable ? "[" : "") + - Twine(mismatchingDims.first.dim) + - (mismatchingDims.first.isScalable ? "]" : "") + " vs. " + - (mismatchingDims.second.isScalable ? "[" : "") + - Twine(mismatchingDims.second.dim) + - (mismatchingDims.second.isScalable ? "]" : "") + ")") - .str(); - return emitOpError(msg); + return emitOpError("dimension mismatch (") + << (mismatchingDims.first.isScalable ? "[" : "") + << mismatchingDims.first.dim + << (mismatchingDims.first.isScalable ? "]" : "") << " vs. " + << (mismatchingDims.second.isScalable ? "[" : "") + << mismatchingDims.second.dim + << (mismatchingDims.second.isScalable ? "]" : "") << ")"; } if (res == BroadcastableToResult::SourceTypeNotAVector) return emitOpError("source type is not a vector");