blob: 8a7e2c2078f1ff0d2ed6c36f352c23923fc28b44 [file] [log] [blame]
David Monahanc11ba462020-12-03 11:09:46 +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 "ClBaseWorkload.hpp"
Colm Donelan0c479742021-12-10 12:43:54 +00009#include <armnn/backends/WorkloadData.hpp>
David Monahanc11ba462020-12-03 11:09:46 +000010
11#include "ClWorkloadUtils.hpp"
12
13namespace armnn
14{
15
Teresa Charlin588cbdf2022-01-19 15:55:37 +000016struct ClRankWorkload : public ClBaseWorkload<RankQueueDescriptor>
David Monahanc11ba462020-12-03 11:09:46 +000017{
18public:
Teresa Charlin588cbdf2022-01-19 15:55:37 +000019 using ClBaseWorkload<RankQueueDescriptor>::ClBaseWorkload;
David Monahanc11ba462020-12-03 11:09:46 +000020 virtual void Execute() const override
21 {
22 const ClTensorHandle* clTensorHandle = PolymorphicDowncast<const ClTensorHandle*>(m_Data.m_Inputs[0]);
23 const int32_t rank = static_cast<int32_t>(clTensorHandle->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