From 195361c35c76be466fc3529ccd9901550e4afe41 Mon Sep 17 00:00:00 2001
From: David Monahan <david.monahan@arm.com>
Date: Fri, 20 Nov 2020 09:58:54 +0000
Subject: IVGCVSW-5567 armnn_delegate Reshape operator fails

 * Changes to the reshape parser to more closely match the TfLiteParser
 * Added boilerplate checks to avoid potential segfaults with invalid option data
 * Allow fallback to checking for second input tensor if options with invalid data are passed

Signed-off-by: David Monahan <david.monahan@arm.com>
Change-Id: I75c4683148257afc5ce18fafb0b2ac495c5ba2a0
---
 delegate/src/Redefine.hpp | 103 ++++++++++++++++++++++++++++------------------
 1 file changed, 64 insertions(+), 39 deletions(-)

diff --git a/delegate/src/Redefine.hpp b/delegate/src/Redefine.hpp
index a9c27fb9e6..fb58ffdf70 100644
--- a/delegate/src/Redefine.hpp
+++ b/delegate/src/Redefine.hpp
@@ -92,61 +92,86 @@ TfLiteStatus VisitReshapeOperator(DelegateData& delegateData,
     armnn::ReshapeDescriptor reshapeDesc;
 
     // The new shape can be defined by either a second input tensor or by a builtin option, we need to check for both.
-    if (numInputs == 2)
-    {
-        const TfLiteTensor& tfLiteShapeInputTensor = tfLiteTensors[tfLiteNode->inputs->data[1]];
-        if (IsDynamicTensor(tfLiteShapeInputTensor))
-        {
-            TF_LITE_MAYBE_KERNEL_LOG(tfLiteContext,
-                                     "TfLiteArmnnDelegate: Dynamic input tensors are not supported in "
-                                     "operator #%d node #%d: ",
-                                     operatorCode, nodeIndex);
-            return kTfLiteError;
-        }
+    TfLiteReshapeParams* reshapeOptions = reinterpret_cast<TfLiteReshapeParams*>(tfLiteNode->builtin_data);
+    std::vector<int32_t> targetShape;
+    bool targetShapeFound = false;
 
-        // Get the shape data out of the input tensor
-        std::vector<int32_t> targetShape;
-        auto* shapeTensorDataPtr = tflite::GetTensorData<int32_t>(&tfLiteShapeInputTensor);
-        auto shapeTensorNumValues = tfLiteShapeInputTensor.dims->data[0];
-        for (auto i=0; i < shapeTensorNumValues; ++i)
+    if (reshapeOptions != nullptr)
+    {
+        // Options might be set without valid data. we need to check the dimensions are in a valid range.
+        if (reshapeOptions->num_dimensions > 0 && reshapeOptions->num_dimensions <= 8)
         {
-            targetShape.push_back(*(shapeTensorDataPtr+i));
-        }
-
-        // Use the data to create the required tensor shape.
-        if (CreateOutputTensorShape(inputTensorInfo0, targetShape, reshapeDesc) != kTfLiteOk)
-        {
-            TF_LITE_MAYBE_KERNEL_LOG(tfLiteContext,
-                                     "TfLiteArmnnDelegate: At most one component of shape can be -1 in: "
-                                     "operator #%d node #%d: ",
-                                     operatorCode, nodeIndex);
-            return kTfLiteError;
+            uint64_t elementCounter = 1;
+            for (int i=0; i < reshapeOptions->num_dimensions; ++i)
+            {
+                targetShape.push_back(reshapeOptions->shape[i]);
+                elementCounter = elementCounter * reshapeOptions->shape[i];
+            }
+            // Check the number of elements match, otherwise fall back to using the second input tensor.
+            if (elementCounter == inputTensorInfo0.GetNumElements())
+            {
+                targetShapeFound = true;
+            }
         }
     }
-    else if (tfLiteNode->builtin_data)
+    if (!targetShapeFound)
     {
-        std::vector<int32_t> targetShape;
-        TfLiteReshapeParams* reshapeOptions =
-                    reinterpret_cast<TfLiteReshapeParams*>(tfLiteNode->builtin_data);
-        for (int i=0; i < reshapeOptions->num_dimensions; ++i)
+        if (numInputs == 2)
         {
-            targetShape.push_back(reshapeOptions->shape[i]);
+            const TfLiteTensor& tfLiteShapeInputTensor = tfLiteTensors[tfLiteNode->inputs->data[1]];
+            if (IsDynamicTensor(tfLiteShapeInputTensor))
+            {
+                TF_LITE_MAYBE_KERNEL_LOG(tfLiteContext,
+                                         "TfLiteArmnnDelegate: Dynamic input tensors are not supported in "
+                                         "operator #%d node #%d: ",
+                                         operatorCode, nodeIndex);
+                return kTfLiteError;
+            }
+
+            if (tfLiteShapeInputTensor.dims->size != 1)
+            {
+                TF_LITE_MAYBE_KERNEL_LOG(tfLiteContext,
+                         "TfLiteArmnnDelegate: Target 'shape' input is not a 1D tensor in "
+                         "operator #%d node #%d: ",
+                         operatorCode, nodeIndex);
+                return kTfLiteError;
+            }
+
+            // Get the shape data out of the input tensor
+            auto* shapeTensorDataPtr = tflite::GetTensorData<int32_t>(&tfLiteShapeInputTensor);
+            auto shapeTensorNumValues = tfLiteShapeInputTensor.dims->data[0];
+            for (auto i=0; i < shapeTensorNumValues; ++i)
+            {
+                targetShape.push_back(*(shapeTensorDataPtr+i));
+            }
         }
-        if (CreateOutputTensorShape(inputTensorInfo0, targetShape, reshapeDesc) != kTfLiteOk)
+        else
         {
             TF_LITE_MAYBE_KERNEL_LOG(tfLiteContext,
-                                     "TfLiteArmnnDelegate: At most one component of shape can be -1 in: "
-                                     "operator #%d node #%d: ",
+                                     "Target shape not defined in reshape parameters or input tensor. "
+                                     "At least one method required in operator #%d node #%d: ",
                                      operatorCode, nodeIndex);
             return kTfLiteError;
         }
     }
-    else
+
+    // Use the data to create the required tensor shape.
+    if (CreateOutputTensorShape(inputTensorInfo0, targetShape, reshapeDesc) != kTfLiteOk)
     {
         TF_LITE_MAYBE_KERNEL_LOG(tfLiteContext,
-                                 "Target shape not defined in reshape parameters or input tensor. "
-                                 "At least one method required in operator #%d node #%d: ",
+                                 "TfLiteArmnnDelegate: At most one component of shape can be -1 in: "
+                                 "operator #%d node #%d: ",
                                  operatorCode, nodeIndex);
+        return kTfLiteError;
+    }
+
+    if (reshapeDesc.m_TargetShape.GetNumElements() != inputTensorInfo0.GetNumElements())
+    {
+        TF_LITE_MAYBE_KERNEL_LOG(tfLiteContext,
+                         "TfLiteArmnnDelegate: Reshape, number of elements in output shape does not match input "
+                         "operator #%d node #%d: ",
+                         operatorCode, nodeIndex);
+        return kTfLiteError;
     }
 
     bool isSupported = false;
-- 
cgit v1.2.1