xref: /aosp_15_r20/external/armnn/delegate/classic/src/Transpose.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 <armnn/utility/IgnoreUnused.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 <tensorflow/lite/kernels/internal/tensor_ctypes.h>
15 
16 namespace armnnDelegate
17 {
18 
VisitTransposeOperator(DelegateData & delegateData,TfLiteContext * tfLiteContext,TfLiteNode * tfLiteNode,int nodeIndex,int32_t tfliteTransposeOperatorCode)19 TfLiteStatus VisitTransposeOperator(DelegateData& delegateData,
20                                     TfLiteContext* tfLiteContext,
21                                     TfLiteNode* tfLiteNode,
22                                     int nodeIndex,
23                                     int32_t tfliteTransposeOperatorCode)
24 {
25     TF_LITE_ENSURE_STATUS(ValidateNumInputs(tfLiteContext, tfLiteNode, 2, nodeIndex));
26     TF_LITE_ENSURE_STATUS(ValidateNumOutputs(tfLiteContext, tfLiteNode, 1, nodeIndex));
27 
28     const TfLiteTensor *tfLiteTensors = tfLiteContext->tensors;
29     const TfLiteTensor& tfLiteInputTensor0 = tfLiteTensors[tfLiteNode->inputs->data[0]];
30     if (IsDynamicTensor(tfLiteInputTensor0))
31     {
32         TF_LITE_MAYBE_KERNEL_LOG(tfLiteContext,
33                                  "TfLiteArmnnDelegate: Dynamic input tensors are not supported in "
34                                  "operator #%d node #%d: ",
35                                  tfliteTransposeOperatorCode, nodeIndex);
36 
37         return kTfLiteError;
38     }
39 
40     const TfLiteTensor& tfLiteInputTensor1 = tfLiteTensors[tfLiteNode->inputs->data[1]];
41     if (IsDynamicTensor(tfLiteInputTensor1))
42     {
43         TF_LITE_MAYBE_KERNEL_LOG(tfLiteContext,
44                                  "TfLiteArmnnDelegate: Dynamic input tensors are not supported in "
45                                  "operator #%d node #%d: ",
46                                  tfliteTransposeOperatorCode, nodeIndex);
47         return kTfLiteError;
48     }
49 
50     const TfLiteTensor& tfLiteOutputTensor = tfLiteTensors[tfLiteNode->outputs->data[0]];
51     if (IsDynamicTensor(tfLiteOutputTensor))
52     {
53         TF_LITE_MAYBE_KERNEL_LOG(tfLiteContext,
54                                  "TfLiteArmnnDelegate: Dynamic output tensors are not supported in "
55                                  "operator #%d node #%d: ",
56                                  tfliteTransposeOperatorCode, nodeIndex);
57         return kTfLiteError;
58     }
59 
60     const armnn::TensorInfo& inputTensorInfo0 = GetTensorInfoForTfLiteTensor(tfLiteInputTensor0);
61     const armnn::TensorInfo& outputTensorInfo = GetTensorInfoForTfLiteTensor(tfLiteOutputTensor, true);
62 
63     auto* permTensorDataPtr = tflite::GetTensorData<int32_t>(&tfLiteInputTensor1);
64     unsigned int numEl = tfLiteInputTensor1.dims->data[0];
65 
66     ARMNN_ASSERT( numEl <= static_cast<int>(armnn::MaxNumOfTensorDimensions));
67     ARMNN_ASSERT( tfLiteInputTensor1.dims->size == 1); // ensure only single dimension to the permutation tensor
68 
69     armnn::TransposeDescriptor descriptor(armnn::PermutationVector(
70         reinterpret_cast<const armnn::PermutationVector::ValueType *> (permTensorDataPtr),
71         static_cast<armnn::PermutationVector::SizeType>(numEl)));
72 
73     bool isSupported = false;
74     armnn::BackendId setBackend;
75     auto validateFunc = [&](const armnn::TensorInfo& outputTensorInfo, bool& isSupported)
76     {
77         FORWARD_LAYER_SUPPORT_FUNC("TRANSPOSE",
78                                    tfLiteContext,
79                                    IsTransposeSupported,
80                                    delegateData.m_Backends,
81                                    isSupported,
82                                    setBackend,
83                                    inputTensorInfo0,
84                                    outputTensorInfo,
85                                    descriptor);
86     };
87 
88     if (!delegateData.m_Network)
89     {
90         validateFunc(outputTensorInfo, isSupported);
91         return isSupported ? kTfLiteOk : kTfLiteError;
92     }
93 
94     armnn::IConnectableLayer* transposeLayer = delegateData.m_Network->AddTransposeLayer(descriptor);
95     transposeLayer->SetBackendId(setBackend);
96     ARMNN_ASSERT(transposeLayer != nullptr);
97     ARMNN_ASSERT(transposeLayer->GetNumInputSlots() == 1);     // permutation vector given to descriptor object
98 
99     armnn::IOutputSlot& outputSlot = transposeLayer->GetOutputSlot(0);
100     outputSlot.SetTensorInfo(outputTensorInfo);
101 
102     // try to connect the Constant Inputs if there are any
103     if(ProcessInputs(transposeLayer,delegateData, tfLiteContext, tfLiteNode) != kTfLiteOk )
104     {
105         return kTfLiteError;
106     }
107 
108     return Connect(transposeLayer, tfLiteNode, delegateData);
109 }
110 } // namespace armnnDelegate
111