xref: /aosp_15_r20/external/armnn/src/backends/neon/workloads/NeonBatchNormalizationWorkload.hpp (revision 89c4ff92f2867872bb9e2354d150bf0c8c502810)
1 //
2 // Copyright © 2017 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/IFunction.h>
11 #include <arm_compute/runtime/Tensor.h>
12 
13 #include <memory>
14 
15 namespace armnn
16 {
17 
18 arm_compute::Status NeonBatchNormalizationValidate(const TensorInfo& input,
19                                                    const TensorInfo& output,
20                                                    const TensorInfo& mean,
21                                                    const TensorInfo& var,
22                                                    const TensorInfo& beta,
23                                                    const TensorInfo& gamma,
24                                                    const BatchNormalizationDescriptor& descriptor,
25                                                    const ActivationDescriptor* activationDescriptor = nullptr);
26 
27 class NeonBatchNormalizationWorkload : public NeonBaseWorkload<BatchNormalizationQueueDescriptor>
28 {
29 public:
30     NeonBatchNormalizationWorkload(const BatchNormalizationQueueDescriptor& descriptor,
31                                    const WorkloadInfo& info);
32     virtual void Execute() const override;
33 
34 private:
35     std::unique_ptr<arm_compute::IFunction> m_Layer;
36 
37     std::unique_ptr<arm_compute::Tensor> m_Mean;
38     std::unique_ptr<arm_compute::Tensor> m_Variance;
39     std::unique_ptr<arm_compute::Tensor> m_Gamma;
40     std::unique_ptr<arm_compute::Tensor> m_Beta;
41 
42     void FreeUnusedTensors();
43 };
44 
45 } //namespace armnn
46 
47