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

[llvmlistbot at llvm.org]

Author: Longsheng Mou
Date: 2024-08-16T11:19:59+08:00
New Revision: 019fbcc4d728895f0decab13738b5d89d01b1436

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

LOG: [mlir][tosa] Add missing check for new_shape of `tosa.reshape` (#104394)

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

Added: 
    

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

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp b/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp
index d4e49b6e3c044c..c8e2b04eea0e22 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 != ShapedType::kDynamic && 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>