xref: /aosp_15_r20/external/armnn/src/backends/neon/workloads/NeonConvolution2dWorkload.hpp (revision 89c4ff92f2867872bb9e2354d150bf0c8c502810)
1 //
2 // Copyright © 2017,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/IFunction.h>
11 #include <arm_compute/runtime/Tensor.h>
12 #include <arm_compute/runtime/MemoryManagerOnDemand.h>
13 
14 #include <memory>
15 
16 namespace armnn
17 {
18 
19 arm_compute::Status NeonConvolution2dWorkloadValidate(const TensorInfo& input,
20                                                       const TensorInfo& output,
21                                                       const Convolution2dDescriptor& descriptor,
22                                                       const TensorInfo& weights,
23                                                       const Optional<TensorInfo>& biases,
24                                                       bool isFastMathEnabled = false,
25                                                       const ActivationDescriptor* activationDescriptor = nullptr);
26 
27 class NeonConvolution2dWorkload : public NeonBaseWorkload<Convolution2dQueueDescriptor>
28 {
29 public:
30     using BaseWorkload<Convolution2dQueueDescriptor>::m_Data;
31 
32     NeonConvolution2dWorkload(const Convolution2dQueueDescriptor& descriptor,
33                               const WorkloadInfo& info,
34                               std::shared_ptr<arm_compute::MemoryManagerOnDemand>& memoryManager,
35                               const bool isFastMathENabled = false);
36 
37     void Execute() const override;
38 
39     arm_compute::ConvolutionMethod GetConvolutionMethod() const;
40 
41 private:
42     std::unique_ptr<arm_compute::IFunction> m_ConvolutionLayer;
43 
44     mutable std::unique_ptr<arm_compute::Tensor> m_KernelTensor;
45     mutable std::unique_ptr<arm_compute::Tensor> m_BiasTensor;
46     TensorInfo m_KernelTensorInfo;
47     TensorInfo m_BiasTensorInfo;
48     arm_compute::ConvolutionMethod m_ConvolutionMethod;
49     mutable bool prepared = false;
50 };
51 
52 } //namespace armnn
53