[Mlir-commits] [mlir] [mlir][tosa] Add missing check for new_shape of `tosa.reshape` (PR #104394)

[Longsheng Mou]

https://github.com/CoTinker updated https://github.com/llvm/llvm-project/pull/104394

>From 83c9d0d8f254f1f24374a673d87b0d422659383e Mon Sep 17 00:00:00 2001
From: Longsheng Mou <moulongsheng at huawei.com>
Date: Thu, 15 Aug 2024 11:21:15 +0800
Subject: [PATCH] [mlir][tosa] Add missing check for new_shape of
 `tosa.reshape`

This patch adds check for new_shape of `tosa.reshape`. Tensor dimension
with size less than -1 is invalid.
---
 mlir/lib/Dialect/Tosa/IR/TosaOps.cpp | 7 ++++++-
 mlir/test/Dialect/Tosa/invalid.mlir  | 8 ++++++++
 2 files changed, 14 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp b/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp
index d4e49b6e3c044c..5f1beb0e48c9ed 100644
--- a/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp
+++ b/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp
@@ -990,11 +990,16 @@ llvm::LogicalResult tosa::ReshapeOp::verify() {
     return emitOpError() << "new shape does not match result rank";
 
   for (auto [newShapeDim, outputShapeDim] :
-       zip(getNewShape(), outputType.getShape()))
+       zip(getNewShape(), outputType.getShape())) {
     if (newShapeDim != -1 && outputShapeDim != ShapedType::kDynamic &&
         newShapeDim != outputShapeDim)
       return emitOpError() << "new shape is inconsistent with result shape";
 
+    if (newShapeDim < -1)
+      return emitOpError() << "new shape has invalid tensor dimension size "
+                           << newShapeDim;
+  }
+
   if (inputType.hasStaticShape() && outputType.hasStaticShape()) {
     int64_t inputElementsNum = inputType.getNumElements();
     int64_t outputElementsNum = outputType.getNumElements();
diff --git a/mlir/test/Dialect/Tosa/invalid.mlir b/mlir/test/Dialect/Tosa/invalid.mlir
index e723aef3815ce6..806ba22e1bbe8c 100644
--- a/mlir/test/Dialect/Tosa/invalid.mlir
+++ b/mlir/test/Dialect/Tosa/invalid.mlir
@@ -291,6 +291,14 @@ func.func @test_reshape_invalid_placeholders(%arg0 : tensor<?xf32>) -> () {
 
 // -----
 
+func.func @test_reshape_invalid_tensor_dim(%arg0 : tensor<4x?xf32>) -> () {
+  // expected-error at +1 {{'tosa.reshape' op new shape has invalid tensor dimension size -2}}
+  %0 = "tosa.reshape" (%arg0) {new_shape = array<i64: -2, -1>} : (tensor<4x?xf32>) -> tensor<?x4xf32>
+  return
+}
+
+// -----
+
 func.func @test_reverse_axis_out_of_range(%arg0 : tensor<13x21x3xf32>) -> () {
   // expected-error at +1 {{'tosa.reverse' op expect input tensor rank (3) to be larger than reverse axis (5)}}
   %0 = tosa.reverse %arg0 {axis = 5 : i32} : (tensor<13x21x3xf32>) -> tensor<?x?x?xi32>