diff options
author | Omar Al Khatib <omar.alkhatib@arm.com> | 2024-01-02 14:45:07 +0000 |
---|---|---|
committer | Omar Al Khatib <omar.alkhatib@arm.com> | 2024-03-12 15:45:42 +0000 |
commit | 93e743fbe7d52f4c41fcd90762fc38b95be802f7 (patch) | |
tree | d0ded85f3cf08f3aabcac68caee4842f3e94da4a /src/cpu/operators/CpuSoftmax.cpp | |
parent | d0611c10a08a4e4f78885e76856155a1f05e6720 (diff) | |
download | ComputeLibrary-93e743fbe7d52f4c41fcd90762fc38b95be802f7.tar.gz |
Optimize CpuSoftmaxKernel for axis != 0 and neon kernels
Resolves: COMPMID-6501
Signed-off-by: Omar Al Khatib <omar.alkhatib@arm.com>
Change-Id: I0abd3cbb5f861301f407c443988fb7efaa205b5d
Reviewed-on: https://review.mlplatform.org/c/ml/ComputeLibrary/+/11056
Tested-by: Arm Jenkins <bsgcomp@arm.com>
Reviewed-by: Gunes Bayir <gunes.bayir@arm.com>
Comments-Addressed: Arm Jenkins <bsgcomp@arm.com>
Benchmark: Arm Jenkins <bsgcomp@arm.com>
Diffstat (limited to 'src/cpu/operators/CpuSoftmax.cpp')
-rw-r--r-- | src/cpu/operators/CpuSoftmax.cpp | 90 |
1 files changed, 14 insertions, 76 deletions
diff --git a/src/cpu/operators/CpuSoftmax.cpp b/src/cpu/operators/CpuSoftmax.cpp index ae14381ad9..fecee7d765 100644 --- a/src/cpu/operators/CpuSoftmax.cpp +++ b/src/cpu/operators/CpuSoftmax.cpp @@ -1,5 +1,5 @@ /* - * Copyright (c) 2021, 2023 Arm Limited. + * Copyright (c) 2021, 2023-2024 Arm Limited. * * SPDX-License-Identifier: MIT * @@ -41,15 +41,7 @@ namespace arm_compute { namespace cpu { -CpuSoftmaxGeneric::CpuSoftmaxGeneric() - : _permute_input(), - _permute_output(), - _softmax_kernel(), - _tmp(), - _input_permuted(), - _output_permuted(), - _needs_permute(false), - _aux_mem(InternalTensorIdx::COUNT) +CpuSoftmaxGeneric::CpuSoftmaxGeneric() : _softmax_kernel(), _tmp(), _aux_mem(InternalTensorIdx::COUNT) { } @@ -63,17 +55,9 @@ void CpuSoftmaxGeneric::configure(const ITensorInfo *src, ITensorInfo *dst, floa const unsigned int actual_axis = static_cast<unsigned int>(wrap_around(axis, static_cast<int32_t>(src->num_dimensions()))); - _needs_permute = actual_axis > 0; + _axis = actual_axis; - if (_needs_permute) - { - _permute_input.configure(src, &_input_permuted, - softmax_helpers::get_permutation_vector_from_softmax_axis(actual_axis)); - } - - // We want to deal with a 2D input. Either it is the permuted version of the original input (4D case) - // or it is the original input case (2D case) - const ITensorInfo *tmp_input = (_needs_permute ? &_input_permuted : src); + const ITensorInfo *tmp_input = src; TensorInfo tensor_info_tmp; if (is_data_type_quantized_asymmetric(src->data_type())) @@ -88,20 +72,10 @@ void CpuSoftmaxGeneric::configure(const ITensorInfo *src, ITensorInfo *dst, floa // Configure kernels auto sm = std::make_unique<kernels::CpuSoftmaxKernel>(); - if (_needs_permute) - { - // The normalization kernel stores the result in a permuted output tensor - sm->configure(tmp_input, &_output_permuted, beta, is_log, &_tmp); - // Re-permute the permuted output into the requested (4D) output - _permute_output.configure(&_output_permuted, dst, - softmax_helpers::get_permutation_vector_from_softmax_axis(actual_axis)); - } - else - { - // Softmax 2D case - sm->configure(tmp_input, dst, beta, is_log, &_tmp); - } + // Softmax 2D case + sm->configure(tmp_input, dst, beta, is_log, actual_axis, &_tmp); + _softmax_kernel = std::move(sm); if (_tmp.total_size() > 0) @@ -109,11 +83,6 @@ void CpuSoftmaxGeneric::configure(const ITensorInfo *src, ITensorInfo *dst, floa _aux_mem[InternalTensorIdx::TMP] = MemoryInfo(offset_int_vec(InternalTensorIdx::TMP), MemoryLifetime::Temporary, _tmp.total_size()); } - - _aux_mem[InternalTensorIdx::PERMUTED_SRC] = MemoryInfo(offset_int_vec(InternalTensorIdx::PERMUTED_SRC), - MemoryLifetime::Temporary, _input_permuted.total_size()); - _aux_mem[InternalTensorIdx::PERMUTED_DST] = MemoryInfo(offset_int_vec(InternalTensorIdx::PERMUTED_DST), - MemoryLifetime::Temporary, _output_permuted.total_size()); } Status @@ -133,25 +102,11 @@ CpuSoftmaxGeneric::validate(const ITensorInfo *src, const ITensorInfo *dst, floa { tensor_info_tmp = src->clone()->set_data_type(DataType::F32).set_is_resizable(true); } - const unsigned int actual_axis = static_cast<unsigned int>(wrap_around(axis, static_cast<int32_t>(src->num_dimensions()))); - const bool needs_permute = actual_axis > 0; - - if (needs_permute) - { - const PermutationVector permutation_vector = - softmax_helpers::get_permutation_vector_from_softmax_axis(actual_axis); - const TensorShape permuted_shape = - misc::shape_calculator::compute_permutation_output_shape(*src, permutation_vector); - TensorInfo input_permuted(src->clone()->set_tensor_shape(permuted_shape)); - ARM_COMPUTE_RETURN_ON_ERROR(CpuPermute::validate(src, &input_permuted, permutation_vector)); - TensorInfo output_permuted(dst->clone()->set_tensor_shape(permuted_shape)); - ARM_COMPUTE_RETURN_ON_ERROR(CpuPermute::validate(&output_permuted, dst, permutation_vector)); - } - - ARM_COMPUTE_RETURN_ON_ERROR(kernels::CpuSoftmaxKernel::validate(src, dst, beta, is_log, &tensor_info_tmp)); + ARM_COMPUTE_RETURN_ON_ERROR( + kernels::CpuSoftmaxKernel::validate(src, dst, beta, actual_axis, is_log, &tensor_info_tmp)); return Status{}; } @@ -165,34 +120,17 @@ void CpuSoftmaxGeneric::run(ITensorPack &tensors) CpuAuxTensorHandler tmp(offset_int_vec(InternalTensorIdx::TMP), _tmp, tensors, true); - CpuAuxTensorHandler input_permuted(offset_int_vec(InternalTensorIdx::PERMUTED_SRC), _input_permuted, tensors, true); - CpuAuxTensorHandler output_permuted(offset_int_vec(InternalTensorIdx::PERMUTED_DST), _output_permuted, tensors, - true); - ITensorPack softmax_pack; - if (_needs_permute) - { - ITensorPack permute_in_pack = {{TensorType::ACL_SRC, src}, {TensorType::ACL_DST, input_permuted.get()}}; - _permute_input.run(permute_in_pack); + softmax_pack = {{TensorType::ACL_SRC_0, src}, {TensorType::ACL_DST_0, dst}, {TensorType::ACL_DST_1, tmp.get()}}; - softmax_pack = {{TensorType::ACL_SRC_0, input_permuted.get()}, - {TensorType::ACL_DST_0, output_permuted.get()}, - {TensorType::ACL_DST_1, tmp.get()}}; - } - else + if (_axis == 0) { - softmax_pack = {{TensorType::ACL_SRC_0, src}, {TensorType::ACL_DST_0, dst}, {TensorType::ACL_DST_1, tmp.get()}}; + NEScheduler::get().schedule_op(_softmax_kernel.get(), Window::DimY, _softmax_kernel->window(), softmax_pack); } - - NEScheduler::get().schedule_op(_softmax_kernel.get(), Window::DimY, _softmax_kernel->window(), softmax_pack); - - if (_needs_permute) + else { - ITensorPack permute_out_pack; - permute_out_pack.add_tensor(TensorType::ACL_SRC, output_permuted.get()); - permute_out_pack.add_tensor(TensorType::ACL_DST, dst); - _permute_output.run(permute_out_pack); + NEScheduler::get().schedule_op(_softmax_kernel.get(), Window::DimX, _softmax_kernel->window(), softmax_pack); } } |