diff options
Diffstat (limited to 'src/backends/reference/workloads/RefShapeWorkload.hpp')
-rw-r--r-- | src/backends/reference/workloads/RefShapeWorkload.hpp | 48 |
1 files changed, 48 insertions, 0 deletions
diff --git a/src/backends/reference/workloads/RefShapeWorkload.hpp b/src/backends/reference/workloads/RefShapeWorkload.hpp new file mode 100644 index 0000000000..8e2a410b0c --- /dev/null +++ b/src/backends/reference/workloads/RefShapeWorkload.hpp @@ -0,0 +1,48 @@ +// +// Copyright © 2021 Arm Ltd and Contributors. All rights reserved. +// SPDX-License-Identifier: MIT +// + +#pragma once + +#include <backendsCommon/Workload.hpp> +#include <backendsCommon/WorkloadData.hpp> + +#include "RefWorkloadUtils.hpp" + +namespace armnn +{ + +struct RefShapeWorkload : public BaseWorkload<ShapeQueueDescriptor> +{ +public: + using BaseWorkload<ShapeQueueDescriptor>::BaseWorkload; + virtual void Execute() const override + { + Execute(m_Data.m_Inputs, m_Data.m_Outputs); + } + void ExecuteAsync(WorkingMemDescriptor& workingMemDescriptor) override + { + Execute(workingMemDescriptor.m_Inputs, workingMemDescriptor.m_Outputs); + } + +private: + void Execute(std::vector<ITensorHandle*> inputs, std::vector<ITensorHandle*> outputs) const + { + const TensorShape Shape = GetTensorInfo(inputs[0]).GetShape(); + + const TensorInfo& outputInfo = GetTensorInfo(outputs[0]); + + unsigned int numBytes = + GetTensorInfo(inputs[0]).GetNumDimensions() * GetDataTypeSize(outputInfo.GetDataType()); + + std::memcpy(outputs[0]->Map(), &Shape, numBytes); + outputs[0]->Unmap(); + } +}; + +} //namespace armnn + + + + |