1 //
2 // Copyright © 2021,2022-2023 Arm Ltd and Contributors. All rights reserved.
3 // SPDX-License-Identifier: MIT
4 //
5
6 #pragma once
7
8 #include <tensorflow/lite/builtin_ops.h>
9 #include <tensorflow/lite/c/builtin_op_data.h>
10 #include <tensorflow/lite/c/common.h>
11 #include <tensorflow/lite/minimal_logging.h>
12
13 namespace armnnDelegate
14 {
15
VisitPackOperator(DelegateData & delegateData,TfLiteContext * tfLiteContext,TfLiteNode * tfLiteNode,int nodeIndex,int32_t operatorCode)16 TfLiteStatus VisitPackOperator(DelegateData& delegateData,
17 TfLiteContext* tfLiteContext,
18 TfLiteNode* tfLiteNode,
19 int nodeIndex,
20 int32_t operatorCode)
21 {
22 unsigned int numInputs = tfLiteNode->inputs->size;
23 if (numInputs < 1)
24 {
25 TF_LITE_MAYBE_KERNEL_LOG(
26 tfLiteContext, "TfLiteArmnnDelegate: Must have at least one input in (%d != %d) in node #%d",
27 1, numInputs, nodeIndex);
28 return kTfLiteError;
29 }
30
31 TF_LITE_ENSURE_STATUS(ValidateNumOutputs(tfLiteContext, tfLiteNode, 1, nodeIndex));
32
33 const TfLiteTensor* tfLiteTensors = tfLiteContext->tensors;
34
35 // Validate all inputs and get TensorInfo
36 std::vector<armnn::TensorInfo> inputTensorInfos;
37 for (unsigned int i = 0; i < numInputs; ++i)
38 {
39 const TfLiteTensor& tfLiteInputTensor = tfLiteTensors[tfLiteNode->inputs->data[i]];
40 if (!IsValid(tfLiteContext, tfLiteInputTensor, operatorCode, nodeIndex))
41 {
42 return kTfLiteError;
43 }
44
45 armnn::TensorInfo inputTensorInfo = GetTensorInfoForTfLiteTensor(tfLiteInputTensor);
46 inputTensorInfos.emplace_back(inputTensorInfo);
47 }
48
49 // Convert input tensors to const armnn::TensorInfo* type for FORWARD_LAYER_SUPPORT_FUNC.
50 std::vector<const armnn::TensorInfo*> inputConstTensorInfos;
51 std::transform(inputTensorInfos.begin(),
52 inputTensorInfos.end(),
53 std::back_inserter(inputConstTensorInfos),
54 [](armnn::TensorInfo& t)->const armnn::TensorInfo*{ return &t; });
55
56 // Validate output and get TensorInfo
57 const TfLiteTensor& tfLiteOutputTensor = tfLiteTensors[tfLiteNode->outputs->data[0]];
58 if (!IsValid(tfLiteContext, tfLiteOutputTensor, operatorCode, nodeIndex))
59 {
60 return kTfLiteError;
61 }
62
63 const armnn::TensorInfo& outputTensorInfo = GetTensorInfoForTfLiteTensor(tfLiteOutputTensor, true);
64
65 armnn::StackDescriptor desc;
66 desc.m_NumInputs = static_cast<uint32_t>(numInputs);
67
68 // Get axis from TfLite parameters
69 auto* params = reinterpret_cast<TfLitePackParams*>(tfLiteNode->builtin_data);
70 desc.m_Axis = static_cast<uint32_t>(params->axis);
71
72 // Use the tensor shape of the first input as the "correct" input shape in the descriptor
73 desc.m_InputShape = inputTensorInfos[0].GetShape();
74
75 // Check if supported
76 bool isSupported = false;
77 armnn::BackendId setBackend;
78 auto validateFunc = [&](const armnn::TensorInfo& outputTensorInfo, bool& isSupported)
79 {
80 FORWARD_LAYER_SUPPORT_FUNC("STACK",
81 tfLiteContext,
82 IsStackSupported,
83 delegateData.m_Backends,
84 isSupported,
85 setBackend,
86 inputConstTensorInfos,
87 outputTensorInfo,
88 desc);
89 };
90
91 // If the m_Network is a nullptr, this signals that a prerequisite TfLite callback is required to clarify the
92 // support for the operator
93 // If supported, VisitPackOperator will be called again to add the layer to the network as seen below
94 if (!delegateData.m_Network)
95 {
96 validateFunc(outputTensorInfo, isSupported);
97 return isSupported ? kTfLiteOk : kTfLiteError;
98 }
99
100 // The TfLite Pack operator is equivalent to the ArmNN Stack operator
101 armnn::IConnectableLayer* layer = delegateData.m_Network->AddStackLayer(desc);
102 layer->SetBackendId(setBackend);
103 ARMNN_ASSERT(layer != nullptr);
104
105 // Connect the Constant Inputs
106 auto inputsTensorsProcess = ProcessInputs(layer,
107 delegateData,
108 tfLiteContext,
109 tfLiteNode);
110 if (inputsTensorsProcess == kTfLiteError)
111 {
112 return inputsTensorsProcess;
113 }
114
115 armnn::IOutputSlot& outputSlot = layer->GetOutputSlot(0);
116 outputSlot.SetTensorInfo(outputTensorInfo);
117
118 // Connect
119 return Connect(layer, tfLiteNode, delegateData);
120 }
121
122 } // namespace armnnDelegate
123