xref: /aosp_15_r20/external/armnn/src/backends/neon/workloads/NeonRankWorkload.hpp (revision 89c4ff92f2867872bb9e2354d150bf0c8c502810)
1 //
2 // Copyright © 2020 Arm Ltd and Contributors. All rights reserved.
3 // SPDX-License-Identifier: MIT
4 //
5 
6 #pragma once
7 
8 #include "NeonBaseWorkload.hpp"
9 #include <armnn/backends/WorkloadData.hpp>
10 
11 #include "NeonWorkloadUtils.hpp"
12 
13 namespace armnn
14 {
15 
16 struct NeonRankWorkload : public NeonBaseWorkload<RankQueueDescriptor>
17 {
18 public:
19     using NeonBaseWorkload<RankQueueDescriptor>::NeonBaseWorkload;
Executearmnn::NeonRankWorkload20     virtual void Execute() const override
21     {
22         const NeonTensorHandle* neonTensorHandle = PolymorphicDowncast<const NeonTensorHandle*>(m_Data.m_Inputs[0]);
23         const int32_t rank = static_cast<int32_t>(neonTensorHandle->GetShape().GetNumDimensions());
24 
25         std::memcpy(GetOutputTensorData<void>(0, m_Data), &rank, sizeof(int32_t));
26         m_Data.m_Outputs[0]->Unmap();
27     }
28 };
29 
30 } //namespace armnn
31