IVGCVSW-7591 IVGCVSW-7592 Add Gather and GatherNd to Opaque Delegate

Signed-off-by: Kevin May <kevin.may@arm.com>
Change-Id: Id2b6a4f70b1cb50e5f7f7ab4e30487b3816c9ad4

5 files changed, 248 insertions, 0 deletions

diff --git a/delegate/CMakeLists.txt b/delegate/CMakeLists.txt
index 3c5033af5b..06c317a03d 100644
--- a/delegate/CMakeLists.txt
+++ b/delegate/CMakeLists.txt
@@ -272,6 +272,10 @@ if(BUILD_UNIT_TESTS)
              test/Convolution2dTest.cpp
              test/ConvolutionTestHelper.hpp
              test/DepthwiseConvolution2dTest.cpp
+             test/GatherTest.cpp
+             test/GatherTestHelper.hpp
+             test/GatherNdTest.cpp
+             test/GatherNdTestHelper.hpp
              test/TestUtils.hpp
              test/TestUtils.cpp)
 
diff --git a/delegate/opaque/CMakeLists.txt b/delegate/opaque/CMakeLists.txt
index 156f79b0a9..ee364197e5 100644
--- a/delegate/opaque/CMakeLists.txt
+++ b/delegate/opaque/CMakeLists.txt
@@ -11,6 +11,8 @@ list(APPEND armnnOpaqueDelegateObject_sources
         src/armnn_delegate.cpp
         src/BatchSpace.hpp
         src/Convolution.hpp
+        src/Gather.hpp
+        src/GatherNd.hpp
         src/Redefine.hpp
         src/SharedFunctions.cpp
         src/SharedFunctions.hpp)
diff --git a/delegate/opaque/src/Gather.hpp b/delegate/opaque/src/Gather.hpp
index e16969768e..b9eef2bb8e 100644
--- a/delegate/opaque/src/Gather.hpp
+++ b/delegate/opaque/src/Gather.hpp
@@ -2,3 +2,131 @@
 // Copyright © 2023 Arm Ltd and Contributors. All rights reserved.
 // SPDX-License-Identifier: MIT
 //
+
+#pragma once
+
+#include <OpaqueDelegateUtils.hpp>
+
+#include <algorithm>
+#include <iterator>
+#include <string>
+#include <vector>
+
+namespace armnnOpaqueDelegate
+{
+
+    TfLiteStatus VisitGatherOperator(DelegateData& delegateData,
+                                     TfLiteOpaqueContext* tfLiteContext,
+                                     TfLiteOpaqueNode* tfLiteNode,
+                                     int nodeIndex,
+                                     int32_t operatorCode)
+    {
+        TF_LITE_ENSURE_STATUS(ValidateNumInputs(tfLiteContext, tfLiteNode, 2, nodeIndex));
+        TF_LITE_ENSURE_STATUS(ValidateNumOutputs(tfLiteContext, tfLiteNode, 1, nodeIndex));
+
+        int numInputs = 0;
+        const int* inputTensors;
+        if (TfLiteOpaqueNodeInputs(tfLiteNode, &inputTensors, &numInputs) != kTfLiteOk)
+        {
+            TF_LITE_OPAQUE_MAYBE_KERNEL_LOG(
+                    tfLiteContext,
+                    "TfLiteArmnnOpaqueDelegate: Unable to gather input tensor indices from node #%d: ",
+                    nodeIndex);
+            return kTfLiteError;
+        }
+
+        int numOutputs = 0;
+        const int* outputTensors;
+        if (TfLiteOpaqueNodeOutputs(tfLiteNode, &outputTensors, &numOutputs) != kTfLiteOk)
+        {
+            TF_LITE_OPAQUE_MAYBE_KERNEL_LOG(
+                    tfLiteContext,
+                    "TfLiteArmnnOpaqueDelegate: Unable to gather output tensor indices from node #%d: ",
+                    nodeIndex);
+            return kTfLiteError;
+        }
+
+        const TfLiteOpaqueTensor* tfLiteInputTensor = TfLiteOpaqueContextGetOpaqueTensor(tfLiteContext,
+                                                                                         inputTensors[0]);
+        if (!IsValid(tfLiteContext, tfLiteInputTensor, operatorCode, nodeIndex))
+        {
+            return kTfLiteError;
+        }
+
+        const TfLiteOpaqueTensor* tfLiteIndicesTensor = TfLiteOpaqueContextGetOpaqueTensor(tfLiteContext,
+                                                                                           inputTensors[1]);
+        if (!IsValid(tfLiteContext, tfLiteIndicesTensor, operatorCode, nodeIndex))
+        {
+            return kTfLiteError;
+        }
+
+        const TfLiteOpaqueTensor* tfLiteOutputTensor = TfLiteOpaqueContextGetOpaqueTensor(tfLiteContext,
+                                                                                          outputTensors[0]);
+        if (!IsValid(tfLiteContext, tfLiteOutputTensor, operatorCode, nodeIndex))
+        {
+            return kTfLiteError;
+        }
+        auto* tfLiteNodeParameters = reinterpret_cast<TfLiteGatherParams*>(TfLiteOpaqueNodeGetBuiltinData(tfLiteNode));
+        auto axis = tfLiteNodeParameters->axis;
+
+        const armnn::TensorInfo& inputTensorInfo = GetTensorInfoForTfLiteOpaqueTensor(tfLiteInputTensor);
+        const armnn::TensorInfo& indicesTensorInfo = GetTensorInfoForTfLiteOpaqueTensor(tfLiteIndicesTensor);
+        const armnn::TensorInfo& outputTensorInfo = GetTensorInfoForTfLiteOpaqueTensor(tfLiteOutputTensor, true);
+        armnn::GatherDescriptor gatherDescriptor;
+        gatherDescriptor.m_Axis = axis;
+
+        auto inputDimensions = static_cast<int32_t>(inputTensorInfo.GetNumDimensions());
+        auto indicesDimensions = indicesTensorInfo.GetNumDimensions();
+        auto outputDimensions = outputTensorInfo.GetNumDimensions();
+        if (((axis < -inputDimensions) && (axis < 0)) || ((axis >= inputDimensions) && (axis > 0)))
+        {
+            TF_LITE_OPAQUE_MAYBE_KERNEL_LOG(
+                    tfLiteContext,
+                    "TfLiteArmnnDelegate: Operation has invalid axis: %d. It is out of bounds [-%d, %d))",
+                    axis, inputDimensions, inputDimensions);
+            return kTfLiteError;
+        }
+        if (outputDimensions != static_cast<unsigned int>(inputDimensions) + indicesDimensions - 1)
+        {
+            TF_LITE_OPAQUE_MAYBE_KERNEL_LOG(
+                    tfLiteContext,
+                    "Operation has invalid output dimensions: %d. Output must be an (%d + %d - 1)-D tensor",
+                    outputDimensions, inputDimensions, indicesDimensions);
+            return kTfLiteError;
+        }
+
+        armnn::BackendId setBackend;
+        if (!delegateData.m_Network)
+        {
+            // Check if supported
+            bool isSupported = false;
+            FORWARD_LAYER_OPAQUE_SUPPORT_FUNC("GATHER",
+                                              tfLiteContext,
+                                              IsGatherSupported,
+                                              delegateData.m_Backends,
+                                              isSupported,
+                                              setBackend,
+                                              inputTensorInfo,
+                                              indicesTensorInfo,
+                                              outputTensorInfo,
+                                              gatherDescriptor);
+            return isSupported ? kTfLiteOk : kTfLiteError;
+        }
+
+        armnn::IConnectableLayer* layer = delegateData.m_Network->AddGatherLayer(gatherDescriptor);
+        layer->SetBackendId(setBackend);
+        ARMNN_ASSERT(layer != nullptr);
+        layer->GetOutputSlot(0).SetTensorInfo(outputTensorInfo);
+
+        auto inputsTensorsProcess = ProcessInputs(layer,
+                                                  delegateData,
+                                                  tfLiteContext,
+                                                  tfLiteNode);
+        if (inputsTensorsProcess == kTfLiteError)
+        {
+            return inputsTensorsProcess;
+        }
+
+        return Connect(layer, tfLiteContext, tfLiteNode, delegateData);
+    }
+} // namespace armnnOpaqueDelegate
\ No newline at end of file
diff --git a/delegate/opaque/src/GatherNd.hpp b/delegate/opaque/src/GatherNd.hpp
index e16969768e..a23fa889a3 100644
--- a/delegate/opaque/src/GatherNd.hpp
+++ b/delegate/opaque/src/GatherNd.hpp
@@ -2,3 +2,105 @@
 // Copyright © 2023 Arm Ltd and Contributors. All rights reserved.
 // SPDX-License-Identifier: MIT
 //
+
+#pragma once
+
+#include <OpaqueDelegateUtils.hpp>
+
+#include <algorithm>
+#include <iterator>
+#include <string>
+#include <vector>
+
+namespace armnnOpaqueDelegate
+{
+TfLiteStatus VisitGatherNdOperator(DelegateData& delegateData,
+                                   TfLiteOpaqueContext* tfLiteContext,
+                                   TfLiteOpaqueNode* tfLiteNode,
+                                   int nodeIndex,
+                                   int32_t operatorCode)
+{
+    TF_LITE_ENSURE_STATUS(ValidateNumInputs(tfLiteContext, tfLiteNode, 2, nodeIndex));
+    TF_LITE_ENSURE_STATUS(ValidateNumOutputs(tfLiteContext, tfLiteNode, 1, nodeIndex));
+
+    int numInputs = 0;
+    const int* inputTensors;
+    if (TfLiteOpaqueNodeInputs(tfLiteNode, &inputTensors, &numInputs) != kTfLiteOk)
+    {
+        TF_LITE_OPAQUE_MAYBE_KERNEL_LOG(
+                tfLiteContext,
+                "TfLiteArmnnOpaqueDelegate: Unable to gather input tensor indices from node #%d: ",
+                nodeIndex);
+        return kTfLiteError;
+    }
+
+    int numOutputs = 0;
+    const int* outputTensors;
+    if (TfLiteOpaqueNodeOutputs(tfLiteNode, &outputTensors, &numOutputs) != kTfLiteOk)
+    {
+        TF_LITE_OPAQUE_MAYBE_KERNEL_LOG(
+                tfLiteContext,
+                "TfLiteArmnnOpaqueDelegate: Unable to gather output tensor indices from node #%d: ",
+                nodeIndex);
+        return kTfLiteError;
+    }
+
+    const TfLiteOpaqueTensor* tfLiteInputTensor = TfLiteOpaqueContextGetOpaqueTensor(tfLiteContext,
+                                                                                     inputTensors[0]);
+    if (!IsValid(tfLiteContext, tfLiteInputTensor, operatorCode, nodeIndex))
+    {
+        return kTfLiteError;
+    }
+
+    const TfLiteOpaqueTensor* tfLiteIndicesTensor = TfLiteOpaqueContextGetOpaqueTensor(tfLiteContext,
+                                                                                       inputTensors[1]);
+    if (!IsValid(tfLiteContext, tfLiteIndicesTensor, operatorCode, nodeIndex))
+    {
+        return kTfLiteError;
+    }
+
+    const TfLiteOpaqueTensor* tfLiteOutputTensor = TfLiteOpaqueContextGetOpaqueTensor(tfLiteContext,
+                                                                                      outputTensors[0]);
+    if (!IsValid(tfLiteContext, tfLiteOutputTensor, operatorCode, nodeIndex))
+    {
+        return kTfLiteError;
+    }
+
+    const armnn::TensorInfo& inputTensorInfo = GetTensorInfoForTfLiteOpaqueTensor(tfLiteInputTensor);
+    const armnn::TensorInfo& indicesTensorInfo = GetTensorInfoForTfLiteOpaqueTensor(tfLiteIndicesTensor);
+    const armnn::TensorInfo& outputTensorInfo = GetTensorInfoForTfLiteOpaqueTensor(tfLiteOutputTensor, true);
+
+    armnn::BackendId setBackend;
+    if (!delegateData.m_Network)
+    {
+        // Check if supported
+        bool isSupported = false;
+        FORWARD_LAYER_OPAQUE_SUPPORT_FUNC("GATHER_ND",
+                                          tfLiteContext,
+                                          IsGatherNdSupported,
+                                          delegateData.m_Backends,
+                                          isSupported,
+                                          setBackend,
+                                          inputTensorInfo,
+                                          indicesTensorInfo,
+                                          outputTensorInfo);
+        return isSupported ? kTfLiteOk : kTfLiteError;
+    }
+
+    armnn::IConnectableLayer* layer = delegateData.m_Network->AddGatherNdLayer();
+    layer->SetBackendId(setBackend);
+    ARMNN_ASSERT(layer != nullptr);
+    layer->GetOutputSlot(0).SetTensorInfo(outputTensorInfo);
+
+    auto inputsTensorsProcess = ProcessInputs(layer,
+                                              delegateData,
+                                              tfLiteContext,
+                                              tfLiteNode);
+    if (inputsTensorsProcess == kTfLiteError)
+    {
+        return inputsTensorsProcess;
+    }
+
+    return Connect(layer, tfLiteContext, tfLiteNode, delegateData);
+}
+} // namespace armnnOpaqueDelegate
\ No newline at end of file
diff --git a/delegate/opaque/src/armnn_delegate.cpp b/delegate/opaque/src/armnn_delegate.cpp
index 4d43b9271e..38a67e7c46 100644
--- a/delegate/opaque/src/armnn_delegate.cpp
+++ b/delegate/opaque/src/armnn_delegate.cpp
@@ -670,6 +670,18 @@ TfLiteStatus ArmnnSubgraph::VisitNode(DelegateData& delegateData,
                                            tfLiteNode,
                                            nodeIndex,
                                            kTfLiteBuiltinEqual);
+        case kTfLiteBuiltinGather:
+            return VisitGatherOperator(delegateData,
+                                       tfLiteContext,
+                                       tfLiteNode,
+                                       nodeIndex,
+                                       kTfLiteBuiltinGather);
+        case kTfLiteBuiltinGatherNd:
+            return VisitGatherNdOperator(delegateData,
+                                         tfLiteContext,
+                                         tfLiteNode,
+                                         nodeIndex,
+                                         kTfLiteBuiltinGatherNd);
         case kTfLiteBuiltinGreater:
             return VisitComparisonOperator(delegateData,
                                            tfLiteContext,