-
Notifications
You must be signed in to change notification settings - Fork 13.2k
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
[mlir][vector] Fix parser of vector.transfer_read #133721
base: main
Are you sure you want to change the base?
Conversation
Thank you for submitting a Pull Request (PR) to the LLVM Project! This PR will be automatically labeled and the relevant teams will be notified. If you wish to, you can add reviewers by using the "Reviewers" section on this page. If this is not working for you, it is probably because you do not have write permissions for the repository. In which case you can instead tag reviewers by name in a comment by using If you have received no comments on your PR for a week, you can request a review by "ping"ing the PR by adding a comment “Ping”. The common courtesy "ping" rate is once a week. Please remember that you are asking for valuable time from other developers. If you have further questions, they may be answered by the LLVM GitHub User Guide. You can also ask questions in a comment on this PR, on the LLVM Discord or on the forums. |
@llvm/pr-subscribers-mlir @llvm/pr-subscribers-mlir-vector Author: Cedric (douyixuan) ChangesThis PR adds a check in the parser to prevent a crash when vector.transfer_read fails to create minor identity permutation. map. Full diff: https://github.com/llvm/llvm-project/pull/133721.diff 2 Files Affected:
diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index 5a3983699d5a3..d99c8ef0680f4 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -151,7 +151,7 @@ static bool isSupportedCombiningKind(CombiningKind combiningKind,
}
AffineMap mlir::vector::getTransferMinorIdentityMap(ShapedType shapedType,
- VectorType vectorType) {
+ VectorType vectorType) {
int64_t elementVectorRank = 0;
VectorType elementVectorType =
llvm::dyn_cast<VectorType>(shapedType.getElementType());
@@ -164,6 +164,9 @@ AffineMap mlir::vector::getTransferMinorIdentityMap(ShapedType shapedType,
return AffineMap::get(
/*numDims=*/0, /*numSymbols=*/0,
getAffineConstantExpr(0, shapedType.getContext()));
+ if (shapedType.getRank() < vectorType.getRank() - elementVectorRank) {
+ return AffineMap::get(shapedType.getContext());
+ }
return AffineMap::getMinorIdentityMap(
shapedType.getRank(), vectorType.getRank() - elementVectorRank,
shapedType.getContext());
@@ -4260,6 +4263,10 @@ ParseResult TransferReadOp::parse(OpAsmParser &parser, OperationState &result) {
AffineMap permMap;
if (!permMapAttr) {
permMap = getTransferMinorIdentityMap(shapedType, vectorType);
+ if (permMap.isEmpty()) {
+ return parser.emitError(typesLoc,
+ "failed to create minor identity permutation map");
+ }
result.attributes.set(permMapAttrName, AffineMapAttr::get(permMap));
} else {
permMap = llvm::cast<AffineMapAttr>(permMapAttr).getValue();
diff --git a/mlir/test/Dialect/Vector/invalid.mlir b/mlir/test/Dialect/Vector/invalid.mlir
index ea6d0021391fb..667e1615212f4 100644
--- a/mlir/test/Dialect/Vector/invalid.mlir
+++ b/mlir/test/Dialect/Vector/invalid.mlir
@@ -525,6 +525,15 @@ func.func @test_vector.transfer_read(%arg0: memref<?x?xvector<2x3xf32>>) {
// -----
+func.func @test_vector.transfer_read(%arg1: memref<?xindex>) -> vector<3x4xi32> {
+ %c3_i32 = arith.constant 3 : i32
+ // expected-error@+1 {{failed to create minor identity permutation map}}
+ %0 = vector.transfer_read %arg1[%c3_i32, %c3_i32], %c3_i32 : memref<?xindex>, vector<3x4xi32>
+ return %0 : vector<3x4xi32>
+}
+
+// -----
+
func.func @test_vector.transfer_write(%arg0: memref<?x?xf32>) {
%c3 = arith.constant 3 : index
%cst = arith.constant 3.0 : f32
|
@@ -164,6 +164,10 @@ AffineMap mlir::vector::getTransferMinorIdentityMap(ShapedType shapedType, | |||
return AffineMap::get( | |||
/*numDims=*/0, /*numSymbols=*/0, | |||
getAffineConstantExpr(0, shapedType.getContext())); | |||
if (shapedType.getRank() < vectorType.getRank() - elementVectorRank) { | |||
return AffineMap(); // Not enough dimensions in the shaped type to form a |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
The comments couldn't in same line with code.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Does elementVectorRank
matter here? If yes, could you write a test to showcase this?
Separately, given that this is effectively checking some basic pre-conditions, could you move it all the way to the top?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Yes, a test would be like this, a memref with a 1D vector, elementVectorRank
is 1
. There are several in test/Dialect/Vector/vector-transfer-to-vector-load-store.mlir
.
func.func @test_vector.transfer_read(%arg1: memref<?xvector<2xindex>>) -> vector<2xindex> {
%c0 = arith.constant 0 : index
%c1 = arith.constant dense<[0, 0]> : vector<2xindex>
%vector = vector.transfer_read %arg1[%c0], %c1 : memref<?xvector<2xindex>>, vector<2xindex>
return %vector : vector<2xindex>
}
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
OK, could you add tests for this in invalid.mlir? Thanks!
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Actually, the provided MLIR is valid, and the elementVectorRank
calculation is tested in test/Dialect/Vector/vector-transfer-to-vector-load-store.mlir
. Additionally, it is verified by the test case related to this issue, where the elementVectorRank
is 0
. I decided not to include this test this time.
module {
func.func @test_vector.transfer_read(%arg1: memref<?xindex>) -> vector<3x4xi32> {
%c3_i32 = arith.constant 3 : i32
%0 = vector.transfer_read %arg1[%c3_i32, %c3_i32], %c3_i32 : memref<?xindex>, vector<3x4xi32>
return %0 : vector<3x4xi32>
}
}
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thanks! I've left some suggestions.
@@ -164,6 +164,10 @@ AffineMap mlir::vector::getTransferMinorIdentityMap(ShapedType shapedType, | |||
return AffineMap::get( | |||
/*numDims=*/0, /*numSymbols=*/0, | |||
getAffineConstantExpr(0, shapedType.getContext())); | |||
if (shapedType.getRank() < vectorType.getRank() - elementVectorRank) { | |||
return AffineMap(); // Not enough dimensions in the shaped type to form a |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Does elementVectorRank
matter here? If yes, could you write a test to showcase this?
Separately, given that this is effectively checking some basic pre-conditions, could you move it all the way to the top?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thanks for fixing this! LGTM
if (!permMap) { | ||
return parser.emitError( | ||
typesLoc, "failed to create a minor identity map, source rank is less than required for vector rank"); |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
OK, based on your reply elementVectorRank
, this would make more sense to me (apologies about formatting):
if (!permMap) { | |
return parser.emitError( | |
typesLoc, "failed to create a minor identity map, source rank is less than required for vector rank"); | |
if (!permMap) { | |
int64_t elementVectorRank = 0; | |
VectorType elementVectorType = llvm::dyn_cast<VectorType>(shapedType.getElementType()); | |
if (elementVectorType) | |
elementVectorRank += elementVectorType.getRank(); | |
if (shapedType.getRank() < vectorType.getRank() - elementVectorRank) | |
return parser.emitError( | |
typesLoc, "expecting a custom permutation_map when source rank is less than required for vector rank"); |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Initially, I removed the elementVectorRank
calculation since the snippet was repetitive. I noticed there might be other instances where an invalid permMap
is returned from getPermutationMapAttrName
. Following your suggestion, I added the snippet, which seems more reasonable now.
@banach-space Could you please help review this PR?
✅ With the latest revision this PR passed the C/C++ code formatter. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thanks for the updates, just a few final nits.
"expected a custom permutation_map when source " | ||
"rank is less than required for vector rank"); |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
"expected a custom permutation_map when source " | |
"rank is less than required for vector rank"); | |
"expected a custom permutation_map when rank(source) != rank (destination)"; |
Same comment for the transfer_write
part. This way these conditions are identical and can later be extracted to a helper function (I can do it in a separate PR).
@@ -525,6 +525,15 @@ func.func @test_vector.transfer_read(%arg0: memref<?x?xvector<2x3xf32>>) { | |||
|
|||
// ----- | |||
|
|||
func.func @test_vector.transfer_read(%arg1: memref<?xindex>) -> vector<3x4xi32> { |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Let's not mix types: index
-> i32
. Same comment for @test_vector.transfer_write
.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
The test case associated with this issue utilizes i32
instead of index
, while the transfer_write test case serves as its inverse. Should I include an additional test case specifically for i32
?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
From the point of view of the issue that you are fixing, the underlying element type is not relevant. Hence, this should be totally sufficient:
func.func @test_vector.transfer_read(%arg1: memref<?xi32>) -> vector<3x4xi32> {
%c3 = arith.constant 3 : index
// expected-error@+1 {{expected a custom permutation_map when rank(source) != rank(destination)}}
%0 = vector.transfer_read %arg1[%c3, %c3], %c3 : memref<?xi32>, vector<3x4xi32>
return %0 : vector<3x4xi32>
}
Should I include an additional test case specifically for i32?
One test case per Op (using either i32
or index
for all element types) will be sufficient.
Thanks for asking!
This PR adds a check in the parser to prevent a crash when vector.transfer_read fails to create minor identity permutation. map.
Fixes #132851
a.mlir