IVGCVSW-7749 DTS: Fix reshape floating point exception

  * Updated Opaque Delegate, TfliteParser, OnnxParser, and Deserializer to handle the Zero In Shape edge case

Signed-off-by: Tianle Cheng <tianle.cheng@arm.com>
Change-Id: I4a0d1e72a66de1fa56de99af9b6730a84e0ff596

3 files changed, 55 insertions, 5 deletions

diff --git a/src/armnnDeserializer/Deserializer.cpp b/src/armnnDeserializer/Deserializer.cpp
index 6e2c07bf37..f455b1af0a 100644
--- a/src/armnnDeserializer/Deserializer.cpp
+++ b/src/armnnDeserializer/Deserializer.cpp
@@ -2632,7 +2632,23 @@ armnn::TensorInfo IDeserializer::DeserializerImpl::OutputShapeOfReshape(const ar
                std::accumulate(targetDimsIn.begin(), targetDimsIn.end(), -1, std::multiplies<int32_t>()));
 
         auto stretchIndex = static_cast<size_t>(std::distance(targetDimsIn.begin(), stretchDim));
-        outputDims[stretchIndex] = inputTensorInfo.GetNumElements() / targetNumElements;
+        if (targetNumElements == 0)
+        {
+            if (inputTensorInfo.GetNumElements() == 0)
+            {
+                outputDims[stretchIndex] = 0;
+            }
+            else
+            {
+                throw ParseException(
+                        fmt::format("Input to reshape is a tensor with elements, but the requested shape has 0. {}",
+                                    CHECK_LOCATION().AsString()));
+            }
+        }
+        else
+        {
+            outputDims[stretchIndex] = inputTensorInfo.GetNumElements() / targetNumElements;
+        }
     }
 
     TensorShape outputShape = TensorShape(static_cast<unsigned int>(outputDims.size()), outputDims.data());
diff --git a/src/armnnOnnxParser/OnnxParser.cpp b/src/armnnOnnxParser/OnnxParser.cpp
index f165df9e14..5ec99ede74 100644
--- a/src/armnnOnnxParser/OnnxParser.cpp
+++ b/src/armnnOnnxParser/OnnxParser.cpp
@@ -441,10 +441,26 @@ TensorInfo ComputeReshapeInfo(const TensorShape& targetShapeTensor,
                             CHECK_LOCATION().AsString()));
         }
 
-        auto targetNumElements = armnn::numeric_cast<unsigned int>(std::accumulate(targetDims.begin(), targetDims.end(),
-            -1, std::multiplies<int32_t>()));
+        auto targetNumElements = armnn::numeric_cast<unsigned int>(
+                std::accumulate(targetDims.begin(), targetDims.end(), -1, std::multiplies<int32_t>()));
         auto stretchIndex = static_cast<size_t>(std::distance(targetDims.begin(), stretchDim));
-        outDims[stretchIndex] = inShape.GetNumElements() / targetNumElements;
+        if (targetNumElements == 0)
+        {
+            if (inShape.GetNumElements() == 0)
+            {
+                outDims[stretchIndex] = 0;
+            }
+            else
+            {
+                throw ParseException(
+                        fmt::format("Input to reshape is a tensor with elements, but the requested shape has 0. {}",
+                                    CHECK_LOCATION().AsString()));
+            }
+        }
+        else
+        {
+            outDims[stretchIndex] = inShape.GetNumElements() / targetNumElements;
+        }
     }
     TensorShape outShape = TensorShape{static_cast<unsigned int>(outDims.size()), outDims.data()};
     return TensorInfo(outShape, dataType);
diff --git a/src/armnnTfLiteParser/TfLiteParser.cpp b/src/armnnTfLiteParser/TfLiteParser.cpp
index c2be54f5f5..30875d5650 100644
--- a/src/armnnTfLiteParser/TfLiteParser.cpp
+++ b/src/armnnTfLiteParser/TfLiteParser.cpp
@@ -3252,6 +3252,7 @@ void TfLiteParserImpl::ParseActivation(size_t subgraphIndex, size_t operatorInde
     auto outputTensorIndexes = AsUnsignedVector(GetOutputTensorIds(m_Model, subgraphIndex, operatorIndex));
     RegisterOutputSlots(subgraphIndex, operatorIndex, layer, {outputTensorIndexes[0]});
 }
+
 armnn::TensorInfo TfLiteParserImpl::OutputShapeOfReshape(const armnn::TensorInfo& inputTensorInfo,
                                                          const std::vector<int32_t>& targetDimsIn)
 {
@@ -3271,7 +3272,24 @@ armnn::TensorInfo TfLiteParserImpl::OutputShapeOfReshape(const armnn::TensorInfo
                 std::accumulate(targetDimsIn.begin(), targetDimsIn.end(), -1, std::multiplies<int32_t>()));
 
         auto stretchIndex = static_cast<size_t>(std::distance(targetDimsIn.begin(), stretchDim));
-        outputDims[stretchIndex] = inputTensorInfo.GetNumElements() / targetNumElements;
+
+        if (targetNumElements == 0)
+        {
+            if (inputTensorInfo.GetNumElements() == 0)
+            {
+                outputDims[stretchIndex] = 0;
+            }
+            else
+            {
+                throw ParseException(
+                        fmt::format("Input to reshape is a tensor with elements, but the requested shape has 0. {}",
+                                    CHECK_LOCATION().AsString()));
+            }
+        }
+        else
+        {
+            outputDims[stretchIndex] = inputTensorInfo.GetNumElements() / targetNumElements;
+        }
     }
 
     TensorShape outputShape = TensorShape(static_cast<unsigned int>(outputDims.size()), outputDims.data());