[Mlir-commits] [mlir] [mlir][Transforms] Add missing check in applyPermutation (PR #102099)

[llvmlistbot at llvm.org]

https://github.com/DarshanRamakant updated https://github.com/llvm/llvm-project/pull/102099

>From 16ce9900b54830b91968dca3ca99632cce8c0f82 Mon Sep 17 00:00:00 2001
From: Darshan Bhat <darshanbhatsirsi at gmail.com>
Date: Mon, 5 Aug 2024 20:12:32 +0530
Subject: [PATCH] [mlir][Transforms] Add missing check in applyPermutation

The applyPermutation() utility should make sure
that the permutation numbers are within the size
of the input array. Otherwise it will cause a
cryptic array out of bound assertion later.
---
 mlir/include/mlir/Dialect/Utils/IndexingUtils.h        |  3 +++
 .../TosaToLinalg/tosa-to-linalg-invalidperm.mlir       | 10 ++++++++++
 2 files changed, 13 insertions(+)
 create mode 100644 mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-invalidperm.mlir

diff --git a/mlir/include/mlir/Dialect/Utils/IndexingUtils.h b/mlir/include/mlir/Dialect/Utils/IndexingUtils.h
index 7849782e5442b..99218f491ddef 100644
--- a/mlir/include/mlir/Dialect/Utils/IndexingUtils.h
+++ b/mlir/include/mlir/Dialect/Utils/IndexingUtils.h
@@ -202,6 +202,9 @@ SmallVector<T> applyPermutation(ArrayRef<T> input,
                                 ArrayRef<int64_t> permutation) {
   assert(input.size() == permutation.size() &&
          "expected input rank to equal permutation rank");
+  assert(
+      llvm::all_of(permutation, [&](size_t s) { return s < input.size(); }) &&
+      "permutation must be within input bounds");
   auto permutationRange = llvm::map_range(
       llvm::seq<unsigned>(0, input.size()),
       [&](int64_t idx) -> T { return input[permutation[idx]]; });
diff --git a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-invalidperm.mlir b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-invalidperm.mlir
new file mode 100644
index 0000000000000..85378e55cad77
--- /dev/null
+++ b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-invalidperm.mlir
@@ -0,0 +1,10 @@
+// RUN: not --crash mlir-opt %s --pass-pipeline="builtin.module(func.func(tosa-to-linalg-named))" 2>&1 | FileCheck -check-prefix=CHECK %s
+
+func.func @func1() {
+	%arg0 = tensor.empty() : tensor<3x4x5xi32>
+	%1110 = arith.constant dense<[3, 0, 1]> : tensor<3xi32>
+	%143 = tosa.transpose %arg0, %1110: (tensor<3x4x5xi32>, tensor<3xi32>) -> tensor<3x4x5xi32>
+	return 
+}
+
+// CHECK: permutation must be within input bounds