xref: /aosp_15_r20/external/armnn/delegate/classic/src/Pooling.hpp (revision 89c4ff92f2867872bb9e2354d150bf0c8c502810)
1 //
2 // Copyright © 2022-2023 Arm Ltd and Contributors. All rights reserved.
3 // SPDX-License-Identifier: MIT
4 //
5 
6 #pragma once
7 
8 #include <ClassicDelegateUtils.hpp>
9 
10 #include <tensorflow/lite/builtin_ops.h>
11 #include <tensorflow/lite/c/builtin_op_data.h>
12 #include <tensorflow/lite/c/common.h>
13 #include <tensorflow/lite/minimal_logging.h>
14 #include <flatbuffers/flexbuffers.h>
15 
16 namespace armnnDelegate
17 {
18 
VisitPooling2dOperator(DelegateData & delegateData,TfLiteContext * tfLiteContext,TfLiteNode * tfLiteNode,int nodeIndex,int32_t tfLitePoolingOperatorCode)19 TfLiteStatus VisitPooling2dOperator(DelegateData& delegateData,
20                                     TfLiteContext* tfLiteContext,
21                                     TfLiteNode* tfLiteNode,
22                                     int nodeIndex,
23                                     int32_t tfLitePoolingOperatorCode)
24 {
25     TF_LITE_ENSURE_STATUS(ValidateNumInputs(tfLiteContext, tfLiteNode, 1, nodeIndex));
26     TF_LITE_ENSURE_STATUS(ValidateNumOutputs(tfLiteContext, tfLiteNode, 1, nodeIndex));
27 
28     const TfLiteTensor* tfLiteTensors = tfLiteContext->tensors;
29     const TfLiteTensor& tfLiteInputTensor = tfLiteTensors[tfLiteNode->inputs->data[0]];
30     if (IsDynamicTensor(tfLiteInputTensor))
31     {
32         TF_LITE_MAYBE_KERNEL_LOG(
33             tfLiteContext,
34             "TfLiteArmnnDelegate: Dynamic input tensors are not supported in operator #%d node #%d: ",
35             tfLitePoolingOperatorCode, nodeIndex);
36         return kTfLiteError;
37     }
38 
39     const TfLiteTensor& tfLiteOutputTensor = tfLiteTensors[tfLiteNode->outputs->data[0]];
40     if (IsDynamicTensor(tfLiteOutputTensor))
41     {
42         TF_LITE_MAYBE_KERNEL_LOG(
43             tfLiteContext,
44             "TfLiteArmnnDelegate: Dynamic output tensors are not supported in operator #%d node #%d: ",
45             tfLitePoolingOperatorCode, nodeIndex);
46         return kTfLiteError;
47     }
48 
49     const armnn::TensorInfo& inputTensorInfo = GetTensorInfoForTfLiteTensor(tfLiteInputTensor);
50     const armnn::TensorInfo& outputTensorInfo = GetTensorInfoForTfLiteTensor(tfLiteOutputTensor, true);
51 
52     auto* tfLiteNodeParameters = reinterpret_cast<TfLitePoolParams*>(tfLiteNode->builtin_data);
53     TfLiteFusedActivation activationType = kTfLiteActNone;
54     if (tfLiteNodeParameters)
55     {
56         activationType = tfLiteNodeParameters->activation;
57         TfLiteStatus activationStatus = ValidateFusedActivationOperator(delegateData, tfLiteContext, outputTensorInfo,
58                                                                         outputTensorInfo, activationType);
59         if(activationStatus != kTfLiteOk)
60         {
61             return kTfLiteError;
62         }
63 
64     }
65 
66     armnn::PoolingAlgorithm poolingAlgorithm;
67     switch(tfLitePoolingOperatorCode)
68     {
69         case kTfLiteBuiltinAveragePool2d:
70             poolingAlgorithm = armnn::PoolingAlgorithm::Average;
71             break;
72         case kTfLiteBuiltinL2Pool2d:
73             poolingAlgorithm = armnn::PoolingAlgorithm::L2;
74             break;
75         case kTfLiteBuiltinMaxPool2d:
76             poolingAlgorithm = armnn::PoolingAlgorithm::Max;
77             break;
78         default:
79             return kTfLiteError;
80     }
81 
82     armnn::Pooling2dDescriptor descriptor;
83     descriptor.m_PoolType = poolingAlgorithm;
84 
85     descriptor.m_PoolWidth = tfLiteNodeParameters->filter_width;
86     descriptor.m_PoolHeight = tfLiteNodeParameters->filter_height;
87     descriptor.m_StrideX = tfLiteNodeParameters->stride_width;
88     descriptor.m_StrideY = tfLiteNodeParameters->stride_height;
89     descriptor.m_DataLayout = armnn::DataLayout::NHWC;
90 
91     unsigned int inputHeight = inputTensorInfo.GetShape()[1];
92     unsigned int inputWidth  = inputTensorInfo.GetShape()[2];
93 
94     CalcPadding(inputHeight, descriptor.m_PoolHeight, descriptor.m_StrideY, 1u,
95                 descriptor.m_PadTop, descriptor.m_PadBottom, tfLiteNodeParameters->padding);
96     CalcPadding(inputWidth, descriptor.m_PoolWidth, descriptor.m_StrideX, 1u,
97                 descriptor.m_PadLeft, descriptor.m_PadRight, tfLiteNodeParameters->padding);
98 
99     bool isSupported = false;
100     armnn::BackendId setBackend;
101     auto validateFunc = [&](const armnn::TensorInfo& outputTensorInfo, bool& isSupported)
102     {
103         FORWARD_LAYER_SUPPORT_FUNC("POOLING_2D",
104                                    tfLiteContext,
105                                    IsPooling2dSupported,
106                                    delegateData.m_Backends,
107                                    isSupported,
108                                    setBackend,
109                                    inputTensorInfo,
110                                    outputTensorInfo,
111                                    descriptor);
112     };
113 
114     if (!delegateData.m_Network)
115     {
116         validateFunc(outputTensorInfo, isSupported);
117         return isSupported ? kTfLiteOk : kTfLiteError;
118     }
119 
120     armnn::IConnectableLayer* poolingLayer = delegateData.m_Network->AddPooling2dLayer(descriptor);
121     poolingLayer->SetBackendId(setBackend);
122     ARMNN_ASSERT(poolingLayer != nullptr);
123 
124     armnn::IOutputSlot& outputSlot = poolingLayer->GetOutputSlot(0);
125     outputSlot.SetTensorInfo(outputTensorInfo);
126 
127     // try to connect the Constant Inputs if there are any
128     if(ProcessInputs(poolingLayer,delegateData, tfLiteContext, tfLiteNode) != kTfLiteOk )
129     {
130         return kTfLiteError;
131     }
132 
133     if(Connect(poolingLayer, tfLiteNode, delegateData) != kTfLiteOk)
134     {
135         return kTfLiteError;
136     }
137 
138     // Check and create activation
139     return FusedActivation(tfLiteContext, tfLiteNode, activationType, poolingLayer, 0, delegateData);
140 }
141 
VisitPooling3dOperator(DelegateData & delegateData,TfLiteContext * tfLiteContext,TfLiteNode * tfLiteNode,int nodeIndex,std::string customOperatorName)142 TfLiteStatus VisitPooling3dOperator(DelegateData& delegateData,
143                                     TfLiteContext* tfLiteContext,
144                                     TfLiteNode* tfLiteNode,
145                                     int nodeIndex,
146                                     std::string customOperatorName)
147 {
148     TF_LITE_ENSURE_STATUS(ValidateNumInputs(tfLiteContext, tfLiteNode, 1, nodeIndex));
149     TF_LITE_ENSURE_STATUS(ValidateNumOutputs(tfLiteContext, tfLiteNode, 1, nodeIndex));
150 
151     const TfLiteTensor* tfLiteTensors = tfLiteContext->tensors;
152     const TfLiteTensor& tfLiteInputTensor = tfLiteTensors[tfLiteNode->inputs->data[0]];
153     if (IsDynamicTensor(tfLiteInputTensor))
154     {
155         TF_LITE_MAYBE_KERNEL_LOG(
156             tfLiteContext,
157             "TfLiteArmnnDelegate: Dynamic input tensors are not supported in operator #%d node #%d: ",
158             customOperatorName.c_str(), nodeIndex);
159         return kTfLiteError;
160     }
161 
162     const TfLiteTensor& tfLiteOutputTensor = tfLiteTensors[tfLiteNode->outputs->data[0]];
163     if (IsDynamicTensor(tfLiteOutputTensor))
164     {
165         TF_LITE_MAYBE_KERNEL_LOG(
166             tfLiteContext,
167             "TfLiteArmnnDelegate: Dynamic output tensors are not supported in operator #%d node #%d: ",
168             customOperatorName.c_str(), nodeIndex);
169         return kTfLiteError;
170     }
171     // Set the input and output info
172     const armnn::TensorInfo& inputTensorInfo = GetTensorInfoForTfLiteTensor(tfLiteInputTensor);
173     const armnn::TensorInfo& outputTensorInfo = GetTensorInfoForTfLiteTensor(tfLiteOutputTensor, true);
174 
175     // Custom Operators are defined by the name string associated to the operator. Use this to determine
176     // which pooling algorithm to create the armnn operator with. L2 Pooling3D is unsupported in TfLite.
177     armnn::PoolingAlgorithm poolingAlgorithm;
178     if (customOperatorName == "MaxPool3D")
179     {
180         poolingAlgorithm = armnn::PoolingAlgorithm::Max;
181     }
182     else if (customOperatorName == "AveragePool3D")
183     {
184         poolingAlgorithm = armnn::PoolingAlgorithm::Average;
185     }
186     else
187     {
188         return kTfLiteError;
189     }
190     // Create the armnn pool3d descriptor and set the algorithm parsed above.
191     armnn::Pooling3dDescriptor descriptor;
192     descriptor.m_PoolType = poolingAlgorithm;
193 
194     // custom_initial_data and custom_initial_data_size are void* variables defined in the tflite registration
195     // used to access the custom option buffer for the operator.
196     auto custom_data = tfLiteNode->custom_initial_data;
197     auto custom_data_size = tfLiteNode->custom_initial_data_size;
198     // Reinterpret the void* to a byte buffer to access the options data in the flexbuffers map.
199     const flexbuffers::Map& m = flexbuffers::GetRoot(reinterpret_cast<const uint8_t*>(custom_data),
200                                                      custom_data_size).AsMap();
201     // poolDims is a vector of [ 1, Depth, Height, Width, 1 ]
202     const auto poolDims = m["ksize"].AsTypedVector();
203     descriptor.m_PoolWidth = poolDims[3].AsInt32();
204     descriptor.m_PoolHeight = poolDims[2].AsInt32();
205     descriptor.m_PoolDepth = poolDims[1].AsInt32();
206 
207     // strideDimes is a vector of [ 1, Z, Y, X, 1]
208     const auto strideDims = m["strides"].AsTypedVector();
209     descriptor.m_StrideX = strideDims[3].AsInt32();
210     descriptor.m_StrideY = strideDims[2].AsInt32();
211     descriptor.m_StrideZ = strideDims[1].AsInt32();
212     descriptor.m_DataLayout = armnn::DataLayout::NDHWC;
213 
214     unsigned int inputDepth = inputTensorInfo.GetShape()[1];
215     unsigned int inputHeight = inputTensorInfo.GetShape()[2];
216     unsigned int inputWidth = inputTensorInfo.GetShape()[3];
217 
218     // CalcPadding expects a TfLitePadding type. Parse flexbuffers to extract padding string and create TfLitePadding.
219     std::string paddingStr = m["padding"].AsString().str();
220     TfLitePadding padding;
221     if (paddingStr == "VALID")
222     {
223         padding = kTfLitePaddingValid;
224     }
225     else if (paddingStr == "SAME")
226     {
227         padding = kTfLitePaddingSame;
228     }
229     else
230     {
231         padding = kTfLitePaddingUnknown;
232     }
233     // Calculates padding for each pooling dimension separately
234     CalcPadding(inputHeight, descriptor.m_PoolHeight, descriptor.m_StrideY, 1u,
235                 descriptor.m_PadTop, descriptor.m_PadBottom, padding);
236     CalcPadding(inputWidth, descriptor.m_PoolWidth, descriptor.m_StrideX, 1u,
237                 descriptor.m_PadLeft, descriptor.m_PadRight, padding);
238     CalcPadding(inputDepth, descriptor.m_PoolDepth, descriptor.m_StrideZ, 1u,
239                 descriptor.m_PadFront, descriptor.m_PadBack, padding);
240 
241 
242     // Check activation by parsing the string from the flexbuffer map
243     std::string activationTypeStr = m["activation"].AsString().str();
244     TfLiteFusedActivation activationType = kTfLiteActNone;
245 
246     if (activationTypeStr == "kTfLiteActRelu")
247     {
248         activationType = kTfLiteActRelu;
249     }
250     else if (activationTypeStr == "kTfLiteActReluN1To1")
251     {
252         activationType = kTfLiteActReluN1To1;
253     }
254     else if (activationTypeStr == "kTfLiteActRelu6")
255     {
256         activationType = kTfLiteActRelu6;
257     }
258     else if (activationTypeStr == "kTfLiteActTanh")
259     {
260         activationType = kTfLiteActTanh;
261     }
262     else if (activationTypeStr == "kTfLiteActSignBit")
263     {
264         activationType = kTfLiteActSignBit;
265     }
266     else if (activationTypeStr == "kTfLiteActSigmoid")
267     {
268         activationType = kTfLiteActSigmoid;
269     }
270     else
271     {
272         activationType = kTfLiteActNone;
273     }
274 
275     TfLiteStatus activationStatus = ValidateFusedActivationOperator(delegateData, tfLiteContext, outputTensorInfo,
276                                                                     outputTensorInfo, activationType);
277     if(activationStatus != kTfLiteOk)
278     {
279         return kTfLiteError;
280     }
281 
282 
283     // Validate the output info.
284     bool isSupported = false;
285     armnn::BackendId setBackend;
286     auto validateFunc = [&](const armnn::TensorInfo& outputTensorInfo, bool& isSupported) {
287         FORWARD_LAYER_SUPPORT_FUNC("POOLING_3D",
288                                    tfLiteContext,
289                                    IsPooling3dSupported,
290                                    delegateData.m_Backends,
291                                    isSupported,
292                                    setBackend,
293                                    inputTensorInfo,
294                                    outputTensorInfo,
295                                    descriptor);
296     };
297 
298     if (!delegateData.m_Network)
299     {
300         validateFunc(outputTensorInfo, isSupported);
301         return isSupported ? kTfLiteOk : kTfLiteError;
302     }
303 
304     // Create the Layer
305     armnn::IConnectableLayer* poolingLayer = delegateData.m_Network->AddPooling3dLayer(descriptor);
306     poolingLayer->SetBackendId(setBackend);
307     ARMNN_ASSERT(poolingLayer != nullptr);
308 
309     // Create and set output slots
310     armnn::IOutputSlot& outputSlot = poolingLayer->GetOutputSlot(0);
311     outputSlot.SetTensorInfo(outputTensorInfo);
312 
313     // try to connect the Constant Inputs if there are any
314     if(ProcessInputs(poolingLayer,delegateData, tfLiteContext, tfLiteNode) != kTfLiteOk )
315     {
316         return kTfLiteError;
317     }
318 
319     if(Connect(poolingLayer, tfLiteNode, delegateData) != kTfLiteOk)
320     {
321         return kTfLiteError;
322     }
323 
324     return FusedActivation(tfLiteContext, tfLiteNode, activationType, poolingLayer, 0, delegateData);
325 }
326 
327 } // namespace armnnDelegate
328