From 1b703d758bcabf66b447cdaab7ab2dc43046563c Mon Sep 17 00:00:00 2001
From: Filip Jankovic <Filip.Jankovic@syrmia.com>
Date: Mon, 4 Mar 2024 11:56:06 +0100
Subject: [PATCH] [mlir][tosa] Prevent div by 0 in
 ReshapeOp::inferReturnTypeComponents

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

diff --git a/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp b/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp
index 62d07859e32f6..bcde87ff158b4 100644
--- a/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp
+++ b/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp
@@ -936,8 +936,12 @@ LogicalResult tosa::ReshapeOp::inferReturnTypeComponents(
 
   // Determine the length of the dynamic dimension.
   for (auto &val : newShapeValue) {
-    if (ShapedType::isDynamic(val))
+    if (ShapedType::isDynamic(val)) {
+      if (staticMul == 0) {
+        return failure();
+      }
       val = numElements / staticMul;
+    }
   }
 
   inferredReturnShapes.push_back(