diff options
author | David Monahan <david.monahan@arm.com> | 2020-11-20 09:58:54 +0000 |
---|---|---|
committer | David Monahan <david.monahan@arm.com> | 2020-11-20 09:58:54 +0000 |
commit | e03d9c2efd11090515284cc2067b451affa04151 (patch) | |
tree | 3d858ad6783c4c9623dbefdbe9e9df9fb3e36afb /delegate/src | |
parent | 187b3a73e63f997b2d4215b6fb1477b55c9764f9 (diff) | |
download | armnn-e03d9c2efd11090515284cc2067b451affa04151.tar.gz |
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
Diffstat (limited to 'delegate/src')
-rw-r--r-- | delegate/src/Redefine.hpp | 103 |
1 files 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; |