1 // 2 // Copyright © 2022 Arm Ltd and Contributors. All rights reserved. 3 // SPDX-License-Identifier: MIT 4 // 5 6 #pragma once 7 8 #include "NeonBaseWorkload.hpp" 9 10 #include "arm_compute/runtime/Tensor.h" 11 #include "arm_compute/runtime/NEON/functions/NEGather.h" 12 #include "arm_compute/runtime/NEON/functions/NEPixelWiseMultiplication.h" 13 #include "arm_compute/runtime/NEON/functions/NEReductionOperation.h" 14 #include "arm_compute/runtime/NEON/functions/NEReshapeLayer.h" 15 16 namespace armnn 17 { 18 arm_compute::Status NeonGatherNdWorkloadValidate(const TensorInfo& input, 19 const TensorInfo& indices, 20 const TensorInfo& output); 21 22 class NeonGatherNdWorkload : public NeonBaseWorkload<GatherNdQueueDescriptor> 23 { 24 public: 25 NeonGatherNdWorkload(const GatherNdQueueDescriptor& descriptor, const WorkloadInfo& info); 26 virtual void Execute() const override; 27 28 private: 29 arm_compute::Tensor m_FlattenedCoeff; 30 arm_compute::Tensor m_OutputMul; 31 arm_compute::Tensor m_FlattenedIndices; 32 arm_compute::Tensor m_OutputGather; 33 34 mutable arm_compute::NEPixelWiseMultiplication m_MulLayer; 35 mutable arm_compute::NEReductionOperation m_ReduceSumLayer; 36 mutable arm_compute::NEGather m_GatherLayer; 37 mutable arm_compute::NEReshapeLayer m_ReshapeLayer; 38 39 }; 40 41 } //namespace armnn