[AutoBump] Merge with fixes of e6eb94d3 (Sep 20) (8)

mlir/lib/Dialect/Tosa/IR/TosaOps.cpp

@@ -30,6 +30,8 @@
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/TypeSwitch.h"
 
+#include <numeric>
+
 using namespace mlir;
 using namespace mlir::tosa;
 
@@ -888,7 +890,7 @@ LogicalResult tosa::SliceOp::verify() {
   if (!inputType || !outputType)
     return success();
 
- if (inputType.getRank() != outputType.getRank()) {
+  if (inputType.getRank() != outputType.getRank()) {
     return emitOpError() << "rank of input (" << inputType.getRank()
                          << ") and output (" << outputType.getRank()
                          << ") must match";
@@ -1077,28 +1079,42 @@ llvm::LogicalResult tosa::ReshapeOp::verify() {
                            << newShapeDim;
   }
 
-  if (inputType.hasStaticShape() && outputType.hasStaticShape()) {
+  if (inputType.hasStaticShape()) {
     int64_t inputElementsNum = inputType.getNumElements();
-    int64_t outputElementsNum = outputType.getNumElements();
-    if (inputElementsNum != outputElementsNum) {
-      return emitOpError() << "cannot reshape " << inputElementsNum
-                           << " elements into " << outputElementsNum;
-    }
+    if (outputType.hasStaticShape()) {
+      int64_t outputElementsNum = outputType.getNumElements();
+      if (inputElementsNum != outputElementsNum) {
+        return emitOpError() << "cannot reshape " << inputElementsNum
+                             << " elements into " << outputElementsNum;
+      }
+
+      if ((int64_t)getNewShape().size() != outputType.getRank()) {
+        return emitOpError()
+               << "rank of newShape (" << getNewShape().size()
+               << ") and output (" << outputType.getRank() << ") must match";
+      }
 
-    if ((int64_t)getNewShape().size() != outputType.getRank()) {
-      return emitOpError() << "rank of newShape (" << getNewShape().size()
-                           << ") and output (" << outputType.getRank()
-                           << ") must match";
+      for (int64_t dim = 0; dim < outputType.getRank(); ++dim) {
+        if (getNewShape()[dim] != -1 &&
+            getNewShape()[dim] != outputType.getShape()[dim]) {
+          return emitOpError()
+                 << "newShape attribute (" << getNewShape()[dim]
+                 << ") does not match output type ("
+                 << outputType.getShape()[dim] << ") in dimension " << dim;
+        }
+      }
     }
 
-    for (int64_t dim = 0; dim < outputType.getRank(); ++dim) {
-      if (getNewShape()[dim] != -1 &&
-          getNewShape()[dim] != outputType.getShape()[dim]) {
-        return emitOpError()
-               << "newShape attribute (" << getNewShape()[dim]
-               << ") does not match output type (" << outputType.getShape()[dim]
-               << ") in dimension " << dim;
-      }
+    // AMD: Switched checks with > to >= to allow zero dimensions
+    int64_t newShapeElementsNum = std::accumulate(
+        getNewShape().begin(), getNewShape().end(), 1LL,
+        [](int64_t acc, int64_t dim) { return (dim >= 0) ? acc * dim : acc; });
+    bool isStaticNewShape =
+        llvm::all_of(getNewShape(), [](int64_t s) { return s >= 0; });
+    if ((isStaticNewShape && inputElementsNum != newShapeElementsNum) ||
+        (!isStaticNewShape && newShapeElementsNum > inputElementsNum)) {
+      return emitOpError() << "cannot reshape " << inputElementsNum
+                           << " elements into " << newShapeElementsNum;
     }
   }
 

mlir/test/Dialect/Tosa/invalid.mlir

@@ -404,6 +404,22 @@ func.func @test_reshape_invalid_size(%arg0 : tensor<2x4xf32>) -> () {
 
 // -----
 
+func.func @test_reshape_invalid_newshape(%arg0 : tensor<1xf32>) -> () {
+  // expected-error@+1 {{'tosa.reshape' op cannot reshape 1 elements into 4}}
+  %0 = "tosa.reshape"(%arg0) {new_shape = array<i64: -1, 4>} : (tensor<1xf32>) -> tensor<?x4xf32>
+  return
+}
+
+// -----
+
+func.func @test_reshape_invalid_newshape(%arg0 : tensor<8xf32>) -> () {
+  // expected-error@+1 {{'tosa.reshape' op cannot reshape 8 elements into 4}}
+  %0 = "tosa.reshape"(%arg0) {new_shape = array<i64: 1, 4>} : (tensor<8xf32>) -> tensor<?x4xf32>
+  return
+}
+
+// -----
+
 func.func @test_reshape_invalid_placeholders(%arg0 : tensor<?xf32>) -> () {
   // expected-error@+1 {{'tosa.reshape' op expected at most one target dimension to be -1}}
   %0 = "tosa.reshape"(%arg0) {new_shape = array<i64: 2, -1, -1>} : (tensor<?xf32>) -> tensor<2x?x?xf32>