blob: 22bf7af5e502c50d0003e8b7b21033407bcd1450 [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
Teresa Charlin588cbdf2022-01-19 15:55:37 +00008#include "NeonBaseWorkload.hpp"
Colm Donelan0c479742021-12-10 12:43:54 +00009#include <armnn/backends/WorkloadData.hpp>
David Monahan97451b42020-12-03 09:48:06 +000010
11#include "NeonWorkloadUtils.hpp"
12
13namespace armnn
14{
15
Teresa Charlin588cbdf2022-01-19 15:55:37 +000016struct NeonRankWorkload : public NeonBaseWorkload<RankQueueDescriptor>
David Monahan97451b42020-12-03 09:48:06 +000017{
18public:
Teresa Charlin588cbdf2022-01-19 15:55:37 +000019 using NeonBaseWorkload<RankQueueDescriptor>::NeonBaseWorkload;
David Monahan97451b42020-12-03 09:48:06 +000020 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