xref: /aosp_15_r20/external/armnn/src/backends/neon/workloads/NeonBatchToSpaceNdWorkload.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/NEBatchToSpaceLayer.h>
14 
15 namespace armnn
16 {
17 
18 arm_compute::Status NeonBatchToSpaceNdWorkloadValidate(const TensorInfo& input,
19                                                        const TensorInfo& output,
20                                                        const BatchToSpaceNdDescriptor& descriptor);
21 
22 class NeonBatchToSpaceNdWorkload : public NeonBaseWorkload<BatchToSpaceNdQueueDescriptor>
23 {
24 public:
25     using NeonBaseWorkload<BatchToSpaceNdQueueDescriptor>::NeonBaseWorkload;
26 
27     NeonBatchToSpaceNdWorkload(const BatchToSpaceNdQueueDescriptor& descriptor, const WorkloadInfo& info);
28 
29     virtual void Execute() const override;
30 
31 private:
32     mutable std::unique_ptr<arm_compute::NEBatchToSpaceLayer> m_Layer;
33 };
34 
35 }
36