diff options
Diffstat (limited to 'src/cpu/kernels/CpuSoftmaxKernel.cpp')
| -rw-r--r-- | src/cpu/kernels/CpuSoftmaxKernel.cpp | 248 |
1 files changed, 248 insertions, 0 deletions
diff --git a/src/cpu/kernels/CpuSoftmaxKernel.cpp b/src/cpu/kernels/CpuSoftmaxKernel.cpp new file mode 100644 index 0000000000..5cf81f815c --- /dev/null +++ b/src/cpu/kernels/CpuSoftmaxKernel.cpp @@ -0,0 +1,248 @@ +/* + * Copyright (c) 2017-2024 Arm Limited. + * + * SPDX-License-Identifier: MIT + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to + * deal in the Software without restriction, including without limitation the + * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or + * sell copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in all + * copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + * SOFTWARE. + */ +#include "src/cpu/kernels/CpuSoftmaxKernel.h" + +#include "arm_compute/core/Error.h" +#include "arm_compute/core/Helpers.h" +#include "arm_compute/core/ITensor.h" +#include "arm_compute/core/TensorInfo.h" +#include "arm_compute/core/Utils.h" +#include "arm_compute/core/Validate.h" +#include "arm_compute/core/Window.h" + +#include "src/core/common/Registrars.h" +#include "src/core/CPP/Validate.h" +#include "src/core/helpers/AutoConfiguration.h" +#include "src/core/helpers/Utils.h" +#include "src/core/helpers/WindowHelpers.h" +#include "src/cpu/kernels/softmax/list.h" + +#include <vector> + +namespace arm_compute +{ +namespace cpu +{ +namespace kernels +{ +namespace +{ +/* Softmax */ +static const std::vector<typename CpuSoftmaxKernel::SoftmaxKernel> available_kernels = { + {"sme2_fp32_softmax", + [](const SoftmaxKernelDataTypeISASelectorData &data) + { return (!data.is_log && data.dt == DataType::F32 && data.isa.sme2 && data.axis == 0); }, + REGISTER_FP32_SME2(sme2_fp32_softmax)}, + {"neon_fp32_softmax", + [](const SoftmaxKernelDataTypeISASelectorData &data) { return (!data.is_log && data.dt == DataType::F32); }, + REGISTER_FP32_NEON(neon_fp32_softmax<false>)}, + {"sme2_fp16_softmax", + [](const SoftmaxKernelDataTypeISASelectorData &data) + { return (!data.is_log && data.dt == DataType::F16 && data.isa.sme2 && data.axis == 0); }, + REGISTER_FP16_SME2(sme2_fp16_softmax)}, + {"neon_fp16_softmax", + [](const SoftmaxKernelDataTypeISASelectorData &data) + { return (!data.is_log && data.dt == DataType::F16) && data.isa.fp16; }, + REGISTER_FP16_NEON(neon_fp16_softmax<false>)}, + {"neon_qu8_softmax", + [](const SoftmaxKernelDataTypeISASelectorData &data) { return (!data.is_log && data.dt == DataType::QASYMM8); }, + REGISTER_QASYMM8_NEON(arm_compute::cpu::neon_qasymm8_softmax<false>)}, + {"neon_qs8_softmax", + [](const SoftmaxKernelDataTypeISASelectorData &data) + { return (!data.is_log && data.dt == DataType::QASYMM8_SIGNED); }, + REGISTER_QASYMM8_SIGNED_NEON(arm_compute::cpu::neon_qasymm8_signed_softmax<false>)}, + {"neon_fp32_log_softmax", + [](const SoftmaxKernelDataTypeISASelectorData &data) { return (data.is_log && data.dt == DataType::F32); }, + REGISTER_FP32_NEON(neon_fp32_softmax<true>)}, + {"neon_fp16_log_softmax", + [](const SoftmaxKernelDataTypeISASelectorData &data) + { return (data.is_log && data.dt == DataType::F16) && data.isa.fp16; }, + REGISTER_FP16_NEON(neon_fp16_softmax<true>)}, + {"neon_qu8_log_softmax", + [](const SoftmaxKernelDataTypeISASelectorData &data) { return (data.is_log && data.dt == DataType::QASYMM8); }, + REGISTER_QASYMM8_NEON(arm_compute::cpu::neon_qasymm8_softmax<true>)}, + {"neon_qs8_log_softmax", + [](const SoftmaxKernelDataTypeISASelectorData &data) + { return (data.is_log && data.dt == DataType::QASYMM8_SIGNED); }, + REGISTER_QASYMM8_SIGNED_NEON(arm_compute::cpu::neon_qasymm8_signed_softmax<true>)}, +}; + +Status validate_arguments_softmax( + const ITensorInfo &src, const ITensorInfo &dst, float beta, int axis, const ITensorInfo &tmp, bool is_log) +{ + ARM_COMPUTE_UNUSED(beta); + // Check input + ARM_COMPUTE_RETURN_ERROR_ON_CPU_F16_UNSUPPORTED(&src); + ARM_COMPUTE_RETURN_ERROR_ON_DATA_TYPE_CHANNEL_NOT_IN(&src, 1, DataType::QASYMM8, DataType::QASYMM8_SIGNED, + DataType::F16, DataType::F32); + + ARM_COMPUTE_RETURN_ERROR_ON(axis < 0 || axis > 3); + + const bool is_quantized_asymmetric = is_data_type_quantized_asymmetric(src.data_type()); + + // Check output if configured + if (dst.total_size() != 0) + { + const QuantizationInfo output_quantization = + is_quantized_asymmetric ? arm_compute::get_softmax_output_quantization_info(src.data_type(), is_log) + : dst.quantization_info(); + ARM_COMPUTE_RETURN_ERROR_ON_MISMATCHING_DATA_TYPES(&src, &dst); + ARM_COMPUTE_RETURN_ERROR_ON_MISMATCHING_SHAPES(&src, &dst); + ARM_COMPUTE_RETURN_ERROR_ON(dst.quantization_info() != output_quantization); + } + + // Check tmp if configured + if (tmp.total_size() != 0) + { + // We have temporary storage only if src data type is quantized. + // Therefore, tmp data type must be F32 + ARM_COMPUTE_RETURN_ERROR_ON(tmp.data_type() != DataType::F32); + ARM_COMPUTE_RETURN_ERROR_ON(!is_quantized_asymmetric); + + // We could potentially reduce tmp memory if we could predict or make an assumption + // on the maximum number of threads that will run in parallel. + ARM_COMPUTE_RETURN_ERROR_ON_MISMATCHING_SHAPES(&src, &tmp); + } + + return Status{}; +} +} // namespace + +const std::vector<typename CpuSoftmaxKernel::SoftmaxKernel> &CpuSoftmaxKernel::get_available_kernels() +{ + return available_kernels; +} + +void CpuSoftmaxKernel::configure( + const ITensorInfo *src, ITensorInfo *dst, float beta, bool is_log, int axis, ITensorInfo *tmp) +{ + _axis = axis; + + ARM_COMPUTE_ERROR_ON_NULLPTR(src, dst, tmp); + ARM_COMPUTE_ERROR_THROW_ON(validate_arguments_softmax(*src, *dst, beta, axis, *tmp, is_log)); + + // Configure kernel window + const bool is_quantized_asymmetric = is_data_type_quantized_asymmetric(src->data_type()); + + // Output auto initialization if not yet initialized + const QuantizationInfo output_quantization = + is_quantized_asymmetric ? arm_compute::get_softmax_output_quantization_info(src->data_type(), is_log) + : dst->quantization_info(); + auto_init_if_empty(*dst, TensorInfo(*src).set_quantization_info(output_quantization).reset_padding()); + + // Tmp auto initialization if not yet initialized and src is quantized + if (is_quantized_asymmetric) + { + auto_init_if_empty(*tmp, TensorInfo(*src).set_data_type(DataType::F32).reset_padding()); + } + + const auto *uk = CpuSoftmaxKernel::get_implementation( + SoftmaxKernelDataTypeISASelectorData{src->data_type(), CPUInfo::get().get_isa(), is_log, axis}); + ARM_COMPUTE_ERROR_ON(uk == nullptr || uk->ukernel == nullptr); + + std::string kernel_name = is_log ? std::string("CpuLogSoftmaxKernel") : std::string("CpuSoftmaxKernel"); + + _beta = beta; + _run_method = uk->ukernel; + _name = kernel_name.append("/").append(uk->name); + + Window win; + + int vec_size = 16 / dst->element_size(); + + if (_axis == 0) + { + win = calculate_max_window(*dst, Steps()); + + /// TODO:Check dimensions > 0 for holes only. For this, we need + /// a utility function checking if there are holes after some dimension. + if (!has_holes(*dst, dst->num_dimensions() - 1)) + { + win = win.collapse(win, Window::DimY); + } + } + else if (_axis > 0 && _axis <= 3) + { + win = calculate_max_window(*dst, Steps(vec_size)); + } + else + { + ARM_COMPUTE_ERROR("Invalid axis"); + } + + win.set(_axis, Window::Dimension(0, 1, 1)); + + ICpuKernel<CpuSoftmaxKernel>::configure(win); +} + +Status CpuSoftmaxKernel::validate( + const ITensorInfo *src, const ITensorInfo *dst, float beta, int axis, bool is_log, const ITensorInfo *tmp) +{ + ARM_COMPUTE_ERROR_ON_NULLPTR(src, dst, tmp); + ARM_COMPUTE_RETURN_ON_ERROR(validate_arguments_softmax(*src, *dst, beta, axis, *tmp, is_log)); + + return Status{}; +} + +void CpuSoftmaxKernel::run_op(ITensorPack &tensors, const Window &window, const ThreadInfo &info) +{ + ARM_COMPUTE_ERROR_ON_UNCONFIGURED_KERNEL(this); + ARM_COMPUTE_ERROR_ON_INVALID_SUBWINDOW(ICpuKernel<CpuSoftmaxKernel>::window(), window); + ARM_COMPUTE_ERROR_ON(_run_method == nullptr); + + const auto src = tensors.get_const_tensor(TensorType::ACL_SRC_0); + auto dst = tensors.get_tensor(TensorType::ACL_DST_0); + + if (is_data_type_quantized_asymmetric(src->info()->data_type())) + { + auto tmp = tensors.get_tensor(TensorType::ACL_DST_1); + unsigned int num_elems_processed_per_iteration; + if (_axis == 0) + { + num_elems_processed_per_iteration = src->info()->valid_region().shape[_axis]; + } + else + { + //16 QASYMM8/QASYMM8_SIGNED elements can fit into the 16-byte vectors. + num_elems_processed_per_iteration = 16; + } + const unsigned int tmp_size_for_thread = tmp->info()->element_size() * num_elems_processed_per_iteration; + + void *tmp_for_thread = tmp->buffer() + (info.thread_id * tmp_size_for_thread); + _run_method(src, tmp_for_thread, dst, _beta, _axis, window); + } + else + { + _run_method(src, nullptr, dst, _beta, _axis, window); + } +} + +const char *CpuSoftmaxKernel::name() const +{ + return _name.c_str(); +} + +} // namespace kernels +} // namespace cpu +} // namespace arm_compute |