[mlir] Fix crash when verifying linalg.transpose (#131733)

Adds checks in `isPermutationVector` for indices that are outside of the
bounds and removes the assert.

Signed-off-by: Ian Wood <[email protected]>

Author: IanWood1

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);

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}}

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">)