[mlir] Fix crash when verifying linalg.transpose #131733
Merged
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
|
@llvm/pr-subscribers-mlir-linalg Author: Ian Wood (IanWood1) ChangesFull diff: https://github.com/llvm/llvm-project/pull/131733.diff 3 Files Affected:
diff --git a/mlir/lib/Dialect/Utils/IndexingUtils.cpp b/mlir/lib/Dialect/Utils/IndexingUtils.cpp
index 108839a4d90e9..d9edabef6693d 100644
--- a/mlir/lib/Dialect/Utils/IndexingUtils.cpp
+++ b/mlir/lib/Dialect/Utils/IndexingUtils.cpp
@@ -220,10 +220,10 @@ bool mlir::isIdentityPermutation(ArrayRef<int64_t> permutation) {
}
bool mlir::isPermutationVector(ArrayRef<int64_t> interchange) {
- assert(llvm::all_of(interchange, [](int64_t s) { return s >= 0; }) &&
- "permutation must be non-negative");
llvm::SmallDenseSet<int64_t, 4> seenVals;
for (auto val : interchange) {
+ if (val < 0 || static_cast<uint64_t>(val) >= interchange.size())
+ return false;
if (seenVals.count(val))
return false;
seenVals.insert(val);
diff --git a/mlir/test/Dialect/Linalg/invalid.mlir b/mlir/test/Dialect/Linalg/invalid.mlir
index f2283db8f89b2..90ceadebbc1fa 100644
--- a/mlir/test/Dialect/Linalg/invalid.mlir
+++ b/mlir/test/Dialect/Linalg/invalid.mlir
@@ -949,6 +949,28 @@ func.func @transpose_invalid_permutation(%input: tensor<16x32x64xf32>,
// -----
+func.func @transpose_out_of_range_permutation(%input: tensor<16x32x64xf32>,
+ %init: tensor<32x64x16xf32>) -> tensor<32x64x16xf32> {
+ // expected-error @+1 {{'linalg.transpose' op permutation is not valid}}
+ %transpose = linalg.transpose
+ ins(%input:tensor<16x32x64xf32>)
+ outs(%init:tensor<32x64x16xf32>)
+ permutation = [1, 2, 3]
+ func.return %transpose : tensor<32x64x16xf32>
+}
+
+// -----
+
+func.func @transpose_negative_permutation(%input: tensor<16x32x64xf32>,
+ %init: tensor<32x64x16xf32>) -> tensor<32x64x16xf32> {
+ // expected-error @+1 {{'linalg.transpose' op permutation is not valid}}
+ %transpose = linalg.transpose
+ ins(%input:tensor<16x32x64xf32>)
+ outs(%init:tensor<32x64x16xf32>)
+ permutation = [1, 2, -1]
+ func.return %transpose : tensor<32x64x16xf32>
+}
+// -----
func.func @transpose_permutated_dims_mismatch(%input: tensor<16x32x64xf32>,
%init: tensor<32x64x16xf32>) -> tensor<32x64x16xf32> {
// expected-error @+1 {{'linalg.transpose' op dim(result, 0) = 32 doesn't match dim(input, permutation[0]) = 16}}
diff --git a/mlir/test/Dialect/Linalg/transform-op-pack.mlir b/mlir/test/Dialect/Linalg/transform-op-pack.mlir
index b3ad73e8df8e7..620a21896b0c1 100644
--- a/mlir/test/Dialect/Linalg/transform-op-pack.mlir
+++ b/mlir/test/Dialect/Linalg/transform-op-pack.mlir
@@ -595,7 +595,7 @@ module attributes {transform.with_named_sequence} {
%unpack = transform.get_consumers_of_result %1[0]
: (!transform.op<"linalg.generic">) -> (!transform.op<"linalg.unpack">)
%2, %pack_2, %unpack_2 =
- // expected-error @below {{invalid outer_perm}}
+ // expected-error @below {{"outer_perm" is not a valid permutation}}
transform.structured.pack_transpose %unpack with_compute_op(%1)
outer_perm = [1]
: (!transform.op<"linalg.unpack">, !transform.op<"linalg.generic">)
@@ -623,7 +623,7 @@ module attributes {transform.with_named_sequence} {
%unpack = transform.get_consumers_of_result %1[0]
: (!transform.op<"linalg.generic">) -> (!transform.op<"linalg.unpack">)
%2, %pack_2, %unpack_2 =
- // expected-error @below {{invalid inner_perm}}
+ // expected-error @below {{"inner_perm" is not a valid permutation}}
transform.structured.pack_transpose %unpack with_compute_op(%1)
inner_perm = [1]
: (!transform.op<"linalg.unpack">, !transform.op<"linalg.generic">)
|
|
@llvm/pr-subscribers-mlir Author: Ian Wood (IanWood1) ChangesFull diff: https://github.com/llvm/llvm-project/pull/131733.diff 3 Files Affected:
diff --git a/mlir/lib/Dialect/Utils/IndexingUtils.cpp b/mlir/lib/Dialect/Utils/IndexingUtils.cpp
index 108839a4d90e9..d9edabef6693d 100644
--- a/mlir/lib/Dialect/Utils/IndexingUtils.cpp
+++ b/mlir/lib/Dialect/Utils/IndexingUtils.cpp
@@ -220,10 +220,10 @@ bool mlir::isIdentityPermutation(ArrayRef<int64_t> permutation) {
}
bool mlir::isPermutationVector(ArrayRef<int64_t> interchange) {
- assert(llvm::all_of(interchange, [](int64_t s) { return s >= 0; }) &&
- "permutation must be non-negative");
llvm::SmallDenseSet<int64_t, 4> seenVals;
for (auto val : interchange) {
+ if (val < 0 || static_cast<uint64_t>(val) >= interchange.size())
+ return false;
if (seenVals.count(val))
return false;
seenVals.insert(val);
diff --git a/mlir/test/Dialect/Linalg/invalid.mlir b/mlir/test/Dialect/Linalg/invalid.mlir
index f2283db8f89b2..90ceadebbc1fa 100644
--- a/mlir/test/Dialect/Linalg/invalid.mlir
+++ b/mlir/test/Dialect/Linalg/invalid.mlir
@@ -949,6 +949,28 @@ func.func @transpose_invalid_permutation(%input: tensor<16x32x64xf32>,
// -----
+func.func @transpose_out_of_range_permutation(%input: tensor<16x32x64xf32>,
+ %init: tensor<32x64x16xf32>) -> tensor<32x64x16xf32> {
+ // expected-error @+1 {{'linalg.transpose' op permutation is not valid}}
+ %transpose = linalg.transpose
+ ins(%input:tensor<16x32x64xf32>)
+ outs(%init:tensor<32x64x16xf32>)
+ permutation = [1, 2, 3]
+ func.return %transpose : tensor<32x64x16xf32>
+}
+
+// -----
+
+func.func @transpose_negative_permutation(%input: tensor<16x32x64xf32>,
+ %init: tensor<32x64x16xf32>) -> tensor<32x64x16xf32> {
+ // expected-error @+1 {{'linalg.transpose' op permutation is not valid}}
+ %transpose = linalg.transpose
+ ins(%input:tensor<16x32x64xf32>)
+ outs(%init:tensor<32x64x16xf32>)
+ permutation = [1, 2, -1]
+ func.return %transpose : tensor<32x64x16xf32>
+}
+// -----
func.func @transpose_permutated_dims_mismatch(%input: tensor<16x32x64xf32>,
%init: tensor<32x64x16xf32>) -> tensor<32x64x16xf32> {
// expected-error @+1 {{'linalg.transpose' op dim(result, 0) = 32 doesn't match dim(input, permutation[0]) = 16}}
diff --git a/mlir/test/Dialect/Linalg/transform-op-pack.mlir b/mlir/test/Dialect/Linalg/transform-op-pack.mlir
index b3ad73e8df8e7..620a21896b0c1 100644
--- a/mlir/test/Dialect/Linalg/transform-op-pack.mlir
+++ b/mlir/test/Dialect/Linalg/transform-op-pack.mlir
@@ -595,7 +595,7 @@ module attributes {transform.with_named_sequence} {
%unpack = transform.get_consumers_of_result %1[0]
: (!transform.op<"linalg.generic">) -> (!transform.op<"linalg.unpack">)
%2, %pack_2, %unpack_2 =
- // expected-error @below {{invalid outer_perm}}
+ // expected-error @below {{"outer_perm" is not a valid permutation}}
transform.structured.pack_transpose %unpack with_compute_op(%1)
outer_perm = [1]
: (!transform.op<"linalg.unpack">, !transform.op<"linalg.generic">)
@@ -623,7 +623,7 @@ module attributes {transform.with_named_sequence} {
%unpack = transform.get_consumers_of_result %1[0]
: (!transform.op<"linalg.generic">) -> (!transform.op<"linalg.unpack">)
%2, %pack_2, %unpack_2 =
- // expected-error @below {{invalid inner_perm}}
+ // expected-error @below {{"inner_perm" is not a valid permutation}}
transform.structured.pack_transpose %unpack with_compute_op(%1)
inner_perm = [1]
: (!transform.op<"linalg.unpack">, !transform.op<"linalg.generic">)
|
MaheshRavishankar
approved these changes
Mar 18, 2025
Signed-off-by: Ian Wood <[email protected]>
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
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.
Adds checks in
isPermutationVectorfor indices that are outside of the bounds and removes the assert.