From 3e8405337e4af2087929fd1d1d31b95a55b50d5b Mon Sep 17 00:00:00 2001
From: Filip Jankovic <Filip.Jankovic@syrmia.com>
Date: Mon, 4 Mar 2024 11:09:11 +0100
Subject: [PATCH] [mlir] Tosa reshape fix

---
 mlir/lib/Dialect/Tosa/IR/TosaOps.cpp | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp b/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp
index 62d07859e32f6..7b6a42e509e01 100644
--- a/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp
+++ b/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp
@@ -928,10 +928,16 @@ LogicalResult tosa::ReshapeOp::inferReturnTypeComponents(
   // dimension.
   int64_t numElements = inputShape.getNumElements();
   int64_t staticMul = 1;
-  for (auto val : newShapeValue) {
+  int i = 0;
+  for (auto &val : newShapeValue) {
+    if (val == 0) {
+      // Swap 0 dim with corresponding input dim
+      val = inputShape.getDimSize(i);
+    }
     if (!ShapedType::isDynamic(val)) {
       staticMul *= val;
     }
+    i++;
   }
 
   // Determine the length of the dynamic dimension.