diff options
Diffstat (limited to 'src/runtime/NEON/functions/NESpaceToBatchLayer.cpp')
| -rw-r--r-- | src/runtime/NEON/functions/NESpaceToBatchLayer.cpp | 20 |
1 files changed, 10 insertions, 10 deletions
diff --git a/src/runtime/NEON/functions/NESpaceToBatchLayer.cpp b/src/runtime/NEON/functions/NESpaceToBatchLayer.cpp index 10b384157d..e8a84246fe 100644 --- a/src/runtime/NEON/functions/NESpaceToBatchLayer.cpp +++ b/src/runtime/NEON/functions/NESpaceToBatchLayer.cpp @@ -1,5 +1,5 @@ /* - * Copyright (c) 2019-2020 Arm Limited. + * Copyright (c) 2019-2021 Arm Limited. * * SPDX-License-Identifier: MIT * @@ -29,7 +29,7 @@ #include "arm_compute/core/Types.h" #include "arm_compute/core/Validate.h" #include "arm_compute/runtime/NEON/NEScheduler.h" -#include "src/core/NEON/kernels/NEMemsetKernel.h" +#include "arm_compute/runtime/NEON/functions/NEFill.h" #include "src/core/NEON/kernels/NESpaceToBatchLayerKernel.h" namespace arm_compute @@ -37,7 +37,7 @@ namespace arm_compute NESpaceToBatchLayer::~NESpaceToBatchLayer() = default; NESpaceToBatchLayer::NESpaceToBatchLayer() - : _space_to_batch_kernel(), _memset_kernel(), _has_padding(false) + : _space_to_batch_kernel(), _fill_f(), _has_padding(false) { } @@ -47,9 +47,9 @@ void NESpaceToBatchLayer::configure(const ITensor *input, const ITensor *block_s if(input->info()->tensor_shape().total_size() != output->info()->tensor_shape().total_size()) { - _has_padding = true; - _memset_kernel = std::make_unique<NEMemsetKernel>(); - _memset_kernel->configure(output, PixelValue(0, input->info()->data_type(), input->info()->quantization_info())); + _has_padding = true; + _fill_f = std::make_unique<NEFill>(); + _fill_f->configure(output, PixelValue(0, input->info()->data_type(), input->info()->quantization_info())); } _space_to_batch_kernel = std::make_unique<NESpaceToBatchLayerKernel>(); _space_to_batch_kernel->configure(input, block_shape, paddings, output); @@ -61,9 +61,9 @@ void NESpaceToBatchLayer::configure(const ITensor *input, const int block_shape_ if(input->info()->tensor_shape().total_size() != output->info()->tensor_shape().total_size()) { - _has_padding = true; - _memset_kernel = std::make_unique<NEMemsetKernel>(); - _memset_kernel->configure(output, PixelValue(0, input->info()->data_type(), input->info()->quantization_info())); + _has_padding = true; + _fill_f = std::make_unique<NEFill>(); + _fill_f->configure(output, PixelValue(0, input->info()->data_type(), input->info()->quantization_info())); } _space_to_batch_kernel = std::make_unique<NESpaceToBatchLayerKernel>(); _space_to_batch_kernel->configure(input, block_shape_x, block_shape_y, padding_left, padding_right, output); @@ -89,7 +89,7 @@ void NESpaceToBatchLayer::run() // Zero out output only if we have paddings if(_has_padding) { - NEScheduler::get().schedule(_memset_kernel.get(), Window::DimY); + _fill_f->run(); } NEScheduler::get().schedule(_space_to_batch_kernel.get(), Window::DimY); } |