diff options
Diffstat (limited to 'src/armnn/backends/NeonWorkloads/NeonConvolution2dBaseWorkload.cpp')
-rw-r--r-- | src/armnn/backends/NeonWorkloads/NeonConvolution2dBaseWorkload.cpp | 36 |
1 files changed, 32 insertions, 4 deletions
diff --git a/src/armnn/backends/NeonWorkloads/NeonConvolution2dBaseWorkload.cpp b/src/armnn/backends/NeonWorkloads/NeonConvolution2dBaseWorkload.cpp index 10c96d82a6..423f02bcb0 100644 --- a/src/armnn/backends/NeonWorkloads/NeonConvolution2dBaseWorkload.cpp +++ b/src/armnn/backends/NeonWorkloads/NeonConvolution2dBaseWorkload.cpp @@ -12,9 +12,38 @@ namespace armnn { +using namespace armcomputetensorutils; + +arm_compute::Status NeonConvolution2dWorkloadValidate(const TensorInfo& input, + const TensorInfo& output, + const Convolution2dDescriptor& descriptor, + const TensorInfo& weights, + const TensorInfo& biases) +{ + const arm_compute::TensorInfo aclInputInfo = BuildArmComputeTensorInfo(input); + const arm_compute::TensorInfo aclOutputInfo = BuildArmComputeTensorInfo(output); + const arm_compute::TensorInfo aclWeightsInfo = BuildArmComputeTensorInfo(weights); + arm_compute::TensorInfo aclBiasesInfo; + arm_compute::TensorInfo *optionalAclBiasesInfo = nullptr; + + if (descriptor.m_BiasEnabled) + { + aclBiasesInfo = BuildArmComputeTensorInfo(biases); + optionalAclBiasesInfo = &aclBiasesInfo; + } + + arm_compute::PadStrideInfo layerInfo = BuildArmComputePadStrideInfo(descriptor); + + return arm_compute::NEConvolutionLayer::validate(&aclInputInfo, + &aclWeightsInfo, + optionalAclBiasesInfo, + &aclOutputInfo, + layerInfo); +} + template<armnn::DataType dataType> NeonConvolution2dBaseWorkload<dataType>::NeonConvolution2dBaseWorkload(const Convolution2dQueueDescriptor& descriptor, - const WorkloadInfo& info) + const WorkloadInfo& info, std::shared_ptr<arm_compute::MemoryManagerOnDemand>& memoryManager) : TypedWorkload<Convolution2dQueueDescriptor, dataType>(descriptor, info) { using arm_compute::NEDirectConvolutionLayer; @@ -50,7 +79,7 @@ NeonConvolution2dBaseWorkload<dataType>::NeonConvolution2dBaseWorkload(const Con if (preferDirectConvolution) { - auto directConvolutionLayer = std::make_unique<arm_compute::NEDirectConvolutionLayer>(); + auto directConvolutionLayer = std::make_unique<arm_compute::NEDirectConvolutionLayer>(memoryManager); directConvolutionLayer->configure(&input, &m_KernelTensor, optionalBiasTensor, @@ -60,7 +89,7 @@ NeonConvolution2dBaseWorkload<dataType>::NeonConvolution2dBaseWorkload(const Con } else { - auto convolutionLayer = std::make_unique<arm_compute::NEConvolutionLayer>(); + auto convolutionLayer = std::make_unique<arm_compute::NEConvolutionLayer>(memoryManager); convolutionLayer->configure(&input, &m_KernelTensor, optionalBiasTensor, @@ -81,4 +110,3 @@ template class NeonConvolution2dBaseWorkload<DataType::QuantisedAsymm8>; } //namespace armnn - |