[mlir][vector] Fix crash in ReorderCastOpsOnBroadcast with non-vector result #170985
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Men-cotton
requested review from
banach-space,
dcaballe and
nicolasvasilache
as code owners
Member
|
@llvm/pr-subscribers-mlir-vector @llvm/pr-subscribers-mlir Author: Men-cotton (Men-cotton) ChangesFixes a crash in Fixes: #126371 Full diff: https://github.com/llvm/llvm-project/pull/170985.diff 2 Files Affected:
diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
index 726da1e9a3d14..ad16b80a732b3 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
@@ -453,6 +453,8 @@ struct ReorderCastOpsOnBroadcast
PatternRewriter &rewriter) const override {
if (op->getNumOperands() != 1)
return failure();
+ if (!isa<VectorType>(op->getResult(0).getType()))
+ return failure();
auto bcastOp = op->getOperand(0).getDefiningOp<vector::BroadcastOp>();
if (!bcastOp)
return failure();
diff --git a/mlir/test/Dialect/Vector/vector-sink.mlir b/mlir/test/Dialect/Vector/vector-sink.mlir
index beaba52af1841..50ff63b44901a 100644
--- a/mlir/test/Dialect/Vector/vector-sink.mlir
+++ b/mlir/test/Dialect/Vector/vector-sink.mlir
@@ -850,3 +850,20 @@ func.func @negative_store_no_single_use(%arg0: memref<?xf32>, %arg1: index, %arg
vector.store %0, %arg0[%arg1] : memref<?xf32>, vector<1xf32>
return %0 : vector<1xf32>
}
+
+// -----
+
+// CHECK-LABEL: func.func @broadcast_cast_non_vector_result
+// CHECK-SAME: (%[[ARG:.*]]: i64)
+// CHECK: %[[BCAST:.*]] = vector.broadcast %[[ARG]] : i64 to vector<26x7xi64>
+// CHECK: %[[CAST:.*]] = builtin.unrealized_conversion_cast %[[BCAST]] : vector<26x7xi64> to !llvm.array<26 x vector<7xi64>>
+// CHECK: return %[[CAST]] : !llvm.array<26 x vector<7xi64>>
+/// This test ensures that the `ReorderCastOpsOnBroadcast` pattern does not
+/// attempt to reorder a cast operation that produces a non-vector result type.
+/// Previously, this would crash because the pattern assumed the result was a
+/// vector type when creating the new inner broadcast.
+func.func @broadcast_cast_non_vector_result(%arg0: i64) -> !llvm.array<26 x vector<7xi64>> {
+ %0 = vector.broadcast %arg0 : i64 to vector<26x7xi64>
+ %1 = builtin.unrealized_conversion_cast %0 : vector<26x7xi64> to !llvm.array<26 x vector<7xi64>>
+ return %1 : !llvm.array<26 x vector<7xi64>>
+}
|
banach-space
approved these changes
Dec 8, 2025
Contributor
banach-space
left a comment
banach-space
left a comment
There was a problem hiding this comment.
Nice, thank you! LGTM % minor comments
| // CHECK: %[[BCAST:.*]] = vector.broadcast %[[ARG]] : i64 to vector<26x7xi64> | ||
| // CHECK: %[[CAST:.*]] = builtin.unrealized_conversion_cast %[[BCAST]] : vector<26x7xi64> to !llvm.array<26 x vector<7xi64>> | ||
| // CHECK: return %[[CAST]] : !llvm.array<26 x vector<7xi64>> | ||
| /// This test ensures that the `ReorderCastOpsOnBroadcast` pattern does not |
Contributor
There was a problem hiding this comment.
Tests for this pattern are located below this block comment:
Please move it accordingly. Thanks!
Comment on lines
863
to
864
| /// Previously, this would crash because the pattern assumed the result was a | ||
| /// vector type when creating the new inner broadcast. |
Contributor
There was a problem hiding this comment.
[nit] "Previously" is a very relative term ("previously to what?"). I would just drop this sentence, it doesn't add any new info.
| /// attempt to reorder a cast operation that produces a non-vector result type. | ||
| /// Previously, this would crash because the pattern assumed the result was a | ||
| /// vector type when creating the new inner broadcast. | ||
| func.func @broadcast_cast_non_vector_result(%arg0: i64) -> !llvm.array<26 x vector<7xi64>> { |
Contributor
There was a problem hiding this comment.
Please follow our naming convention (as per https://mlir.llvm.org/getting_started/TestingGuide/#test-naming-convention)
Suggested change
| func.func @broadcast_cast_non_vector_result(%arg0: i64) -> !llvm.array<26 x vector<7xi64>> { | |
| func.func @negative_broadcast_cast_non_vector_result(%arg0: i64) -> !llvm.array<26 x vector<7xi64>> { |
Contributor
Author
|
@banach-space |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Fixes a crash in
ReorderCastOpsOnBroadcastby ensuring the cast result is aVectorTypebefore applying the pattern.A regression test has been added to mlir/test/Dialect/Vector/vector-sink.mlir.
Fixes: #126371