Method configure

src/cpu/kernels/CpuSoftmaxKernel.cpp:178–261

}

void CpuSoftmaxKernel::configure(
    const ITensorInfo *src, ITensorInfo *dst, float beta, bool is_log, int axis, ITensorInfo *tmp)
{
    ARM_COMPUTE_TRACE_EVENT(ARM_COMPUTE_PROF_CAT_CPU, ARM_COMPUTE_PROF_LVL_CPU, "CpuSoftmaxKernel::configure");
    _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, CPUInfo::get().get_sme2_vector_length_in_bits()});
    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");
    }