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

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-tosa

Author: Longsheng Mou (CoTinker)

<details>
<summary>Changes</summary>

This patch adds check for new_shape of `tosa.reshape`. Tensor dimension with size less than -1 is invalid. Fix #<!-- -->103062.

---
Full diff: https://github.com/llvm/llvm-project/pull/104394.diff


2 Files Affected:

- (modified) mlir/lib/Dialect/Tosa/IR/TosaOps.cpp (+6-1) 
- (modified) mlir/test/Dialect/Tosa/invalid.mlir (+8) 


``````````diff
diff --git a/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp b/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp
index 39ea7a5b61f5ec..3248ef13de40a7 100644
--- a/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp
+++ b/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp
@@ -967,11 +967,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 e1fcf056480083..cb1545f02ef830 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>

``````````

</details>


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