blob: b8b83af88628f0a2855f4b4f206e8eae3d1441a8 [file] [log] [blame]
David Monahan97451b42020-12-03 09:48:06 +00001//
2// Copyright © 2020 Arm Ltd and Contributors. All rights reserved.
3// SPDX-License-Identifier: MIT
4//
5
6#pragma once
7
Colm Donelan0c479742021-12-10 12:43:54 +00008#include <armnn/backends/Workload.hpp>
9#include <armnn/backends/WorkloadData.hpp>
David Monahan97451b42020-12-03 09:48:06 +000010
11#include "NeonWorkloadUtils.hpp"
12
13namespace armnn
14{
15
16struct NeonRankWorkload : public BaseWorkload<RankQueueDescriptor>
17{
18public:
19 using BaseWorkload<RankQueueDescriptor>::BaseWorkload;
20 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));
David Monahanc11ba462020-12-03 11:09:46 +000026 m_Data.m_Outputs[0]->Unmap();
David Monahan97451b42020-12-03 09:48:06 +000027 }
28};
29
30} //namespace armnn