diff options
author | Teresa Charlin <teresa.charlinreyes@arm.com> | 2022-04-25 18:23:41 +0100 |
---|---|---|
committer | TeresaARM <teresa.charlinreyes@arm.com> | 2022-05-04 12:25:18 +0000 |
commit | d5c0ed24ce91ee0da1dcb5858da16f0f8a3d3172 (patch) | |
tree | 4e05c08d5c224e7df9f7ea358e2c598c6be18ff9 /delegate/src/GatherNd.hpp | |
parent | 91a53eab529d88f78572b1155bfd07eb5de141f4 (diff) | |
download | armnn-d5c0ed24ce91ee0da1dcb5858da16f0f8a3d3172.tar.gz |
IVGCVSW-6858 Add GATHERNd Support to the TfLite Delegate
Signed-off-by: Teresa Charlin <teresa.charlinreyes@arm.com>
Change-Id: I56418875b3bb2ae45b5c69bfeaafa1a6126b8085
Diffstat (limited to 'delegate/src/GatherNd.hpp')
-rw-r--r-- | delegate/src/GatherNd.hpp | 81 |
1 files changed, 81 insertions, 0 deletions
diff --git a/delegate/src/GatherNd.hpp b/delegate/src/GatherNd.hpp new file mode 100644 index 0000000000..b2d7a50870 --- /dev/null +++ b/delegate/src/GatherNd.hpp @@ -0,0 +1,81 @@ +// +// Copyright © 2022 Arm Ltd and Contributors. All rights reserved. +// SPDX-License-Identifier: MIT +// + +#pragma once + +#include "DelegateUtils.hpp" +#include <algorithm> +#include <iterator> +#include <string> +#include <vector> + +namespace armnnDelegate +{ +TfLiteStatus VisitGatherNdOperator(DelegateData& delegateData, + TfLiteContext* tfLiteContext, + TfLiteNode* tfLiteNode, + int nodeIndex, + int32_t operatorCode) +{ + TF_LITE_ENSURE_STATUS(ValidateNumInputs(tfLiteContext, tfLiteNode, 2, nodeIndex)); + TF_LITE_ENSURE_STATUS(ValidateNumOutputs(tfLiteContext, tfLiteNode, 1, nodeIndex)); + + const TfLiteTensor* tfLiteTensors = tfLiteContext->tensors; + + const TfLiteTensor& tfLiteInputTensor = tfLiteTensors[tfLiteNode->inputs->data[0]]; + if (!IsValid(tfLiteContext, tfLiteInputTensor, operatorCode, nodeIndex)) + { + return kTfLiteError; + } + + const TfLiteTensor& tfLiteIndicesTensor = tfLiteTensors[tfLiteNode->inputs->data[1]]; + if (!IsValid(tfLiteContext, tfLiteIndicesTensor, operatorCode, nodeIndex)) + { + return kTfLiteError; + } + + const TfLiteTensor& tfLiteOutputTensor = tfLiteTensors[tfLiteNode->outputs->data[0]]; + if (!IsValid(tfLiteContext, tfLiteOutputTensor, operatorCode, nodeIndex)) + { + return kTfLiteError; + } + + const armnn::TensorInfo& inputTensorInfo = GetTensorInfoForTfLiteTensor(tfLiteInputTensor); + const armnn::TensorInfo& indicesTensorInfo = GetTensorInfoForTfLiteTensor(tfLiteIndicesTensor); + const armnn::TensorInfo& outputTensorInfo = GetTensorInfoForTfLiteTensor(tfLiteOutputTensor); + + if (!delegateData.m_Network) + { + // Check if supported + bool isSupported = false; + FORWARD_LAYER_SUPPORT_FUNC("GATHER_ND", + tfLiteContext, + IsGatherNdSupported, + delegateData.m_Backends, + isSupported, + inputTensorInfo, + indicesTensorInfo, + outputTensorInfo); + return isSupported ? kTfLiteOk : kTfLiteError; + } + + armnn::IConnectableLayer* layer = delegateData.m_Network->AddGatherNdLayer(); + ARMNN_ASSERT(layer != nullptr); + layer->GetOutputSlot(0).SetTensorInfo(outputTensorInfo); + + auto inputsTensorsProcess = ProcessInputs(layer, + delegateData, + tfLiteContext, + tfLiteNode); + if (inputsTensorsProcess == kTfLiteError) + { + return inputsTensorsProcess; + } + + Connect(layer, tfLiteNode, delegateData); + + return kTfLiteOk; +} +} // namespace armnnDelegate
\ No newline at end of file |