[Mlir-commits] [mlir] ab7e8b7 - [mlir][tosa] fix a crash when sliceOp has invalid attribute (#68486) (#70063)

[llvmlistbot at llvm.org]

Author: long.chen
Date: 2023-10-25T17:13:37+08:00
New Revision: ab7e8b76366df9a83847fcce8c9c7e01b64c74de

URL: https://github.com/llvm/llvm-project/commit/ab7e8b76366df9a83847fcce8c9c7e01b64c74de
DIFF: https://github.com/llvm/llvm-project/commit/ab7e8b76366df9a83847fcce8c9c7e01b64c74de.diff

LOG: [mlir][tosa] fix a crash when sliceOp has invalid attribute (#68486) (#70063)

add a verifier for tosa::sliceOp

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
    mlir/lib/Dialect/Tosa/IR/TosaOps.cpp
    mlir/test/Dialect/Tosa/invalid.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td b/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
index 901384eae50176b..c0baf478358c132 100644
--- a/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
+++ b/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
@@ -1621,6 +1621,7 @@ def Tosa_SliceOp : Tosa_InferShapedTypeOp<"slice"> {
 
   let hasCanonicalizer = 1;
   let hasFolder = 1;
+  let hasVerifier = 1;
 }
 
 //===----------------------------------------------------------------------===//

diff  --git a/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp b/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp
index 078e50e857fbb2a..c9e64a67302e772 100644
--- a/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp
+++ b/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp
@@ -820,6 +820,22 @@ LogicalResult tosa::SliceOp::inferReturnTypeComponents(
   return success();
 }
 
+LogicalResult tosa::SliceOp::verify() {
+  auto inputType = llvm::dyn_cast<RankedTensorType>(getInput().getType());
+  if (!inputType)
+    return success();
+
+  if (static_cast<size_t>(inputType.getRank()) != getStart().size())
+    return emitOpError(
+        "length of start attribute is not equal rank of input shape");
+
+  if (static_cast<size_t>(inputType.getRank()) != getSize().size())
+    return emitOpError(
+        "length of size attribute is not equal rank of input shape");
+
+  return success();
+}
+
 LogicalResult tosa::TableOp::inferReturnTypeComponents(
     MLIRContext *context, ::std::optional<Location> location,
     TableOp::Adaptor adaptor,

diff  --git a/mlir/test/Dialect/Tosa/invalid.mlir b/mlir/test/Dialect/Tosa/invalid.mlir
index 7a6b507566eb25d..8a290299b925a7c 100644
--- a/mlir/test/Dialect/Tosa/invalid.mlir
+++ b/mlir/test/Dialect/Tosa/invalid.mlir
@@ -303,3 +303,21 @@ func.func @test_variable_write_shape(%arg0: tensor<1x4x8xi32>) -> () {
   tosa.variable.write @stored_var, %arg0 : tensor<1x4x8xi32>
   return
 }
+
+// -----
+
+func.func @test_slice_invalid_start() {
+  %0 = tensor.empty() : tensor<4x31x31xf32>
+  // expected-error at +1 {{'tosa.slice' op length of start attribute is not equal rank of input shape}}
+  %1 = tosa.slice %0 {size = array<i64: 1, 1, 1>, start = array<i64: 1, 1>} : (tensor<4x31x31xf32>) -> tensor<*xf32>
+  return
+}
+
+// -----
+
+func.func @test_slice_invalid_size() {
+  %0 = tensor.empty() : tensor<4x31x31xf32>
+  // expected-error at +1 {{'tosa.slice' op length of size attribute is not equal rank of input shape}}
+  %1 = tosa.slice %0 {size = array<i64: 1>, start = array<i64: 1, 1, 1>} : (tensor<4x31x31xf32>) -> tensor<*xf32>
+  return
+}