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

[llvmlistbot at llvm.org]

https://github.com/lipracer created https://github.com/llvm/llvm-project/pull/70063

None

>From 5ce2584d3c9318dd4c680baff4dc991a66435951 Mon Sep 17 00:00:00 2001
From: lipracer <lipracer at gmail.com>
Date: Wed, 25 Oct 2023 00:46:17 +0800
Subject: [PATCH] [mlir][tosa] fix a crash when sliceOp has invalid attribute
 (#68486)

---
 mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td   |  1 +
 mlir/lib/Dialect/Tosa/IR/TosaOps.cpp           | 16 ++++++++++++++++
 .../Dialect/Tosa/Transforms/TosaValidation.cpp |  2 +-
 mlir/test/Dialect/Tosa/invalid.mlir            | 18 ++++++++++++++++++
 4 files changed, 36 insertions(+), 1 deletion(-)

diff --git a/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td b/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
index 5cc97469d14c314..c0a2f71cfa06a05 100644
--- a/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
+++ b/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
@@ -1614,6 +1614,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 e03904a1611fc42..f459b3b72705860 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/lib/Dialect/Tosa/Transforms/TosaValidation.cpp b/mlir/lib/Dialect/Tosa/Transforms/TosaValidation.cpp
index 8a2254fc24effe2..eb0682563a2623f 100644
--- a/mlir/lib/Dialect/Tosa/Transforms/TosaValidation.cpp
+++ b/mlir/lib/Dialect/Tosa/Transforms/TosaValidation.cpp
@@ -155,7 +155,7 @@ struct TosaValidation : public tosa::impl::TosaValidationBase<TosaValidation> {
 
   bool levelCheckRank(Operation *op, const Value &v,
                       const std::string &checkDesc) {
-    if (ShapedType type = dyn_cast<ShapedType>(v.getType())) {
+    if (auto type = dyn_cast<RankedTensorType>(v.getType())) {
       if (type.getRank() > tosaLevel.MAX_RANK) {
         op->emitOpError() << "failed level check: " << checkDesc;
         return false;
diff --git a/mlir/test/Dialect/Tosa/invalid.mlir b/mlir/test/Dialect/Tosa/invalid.mlir
index 9233662e88db902..4c0703ad6f61a6c 100644
--- a/mlir/test/Dialect/Tosa/invalid.mlir
+++ b/mlir/test/Dialect/Tosa/invalid.mlir
@@ -248,3 +248,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
+}