xref: /aosp_15_r20/external/armnn/src/backends/neon/workloads/NeonSpaceToBatchNdWorkload.hpp (revision 89c4ff92f2867872bb9e2354d150bf0c8c502810)
1 //
2 // Copyright © 2020 Arm Ltd and Contributors. All rights reserved.
3 // SPDX-License-Identifier: MIT
4 //
5 
6 #pragma once
7 
8 #include <armnn/Tensor.hpp>
9 #include <armnn/Descriptors.hpp>
10 
11 #include "NeonBaseWorkload.hpp"
12 
13 #include <arm_compute/runtime/NEON/functions/NESpaceToBatchLayer.h>
14 
15 namespace armnn
16 {
17 
18 arm_compute::Status NeonSpaceToBatchNdWorkloadValidate(const TensorInfo& input,
19                                                        const TensorInfo& output,
20                                                        const SpaceToBatchNdDescriptor& descriptor);
21 
22 class NeonSpaceToBatchNdWorkload : public NeonBaseWorkload<SpaceToBatchNdQueueDescriptor>
23 {
24 public:
25     using NeonBaseWorkload<SpaceToBatchNdQueueDescriptor>::NeonBaseWorkload;
26 
27     NeonSpaceToBatchNdWorkload(const SpaceToBatchNdQueueDescriptor& descriptor, const WorkloadInfo& info);
28 
29     virtual void Execute() const override;
30 
31 private:
32     mutable std::unique_ptr<arm_compute::NESpaceToBatchLayer> m_Layer;
33 };
34 
35 } //namespace armnn