Method configure

src/runtime/NEON/functions/NEFFTConvolutionLayer.cpp:107–276

NEFFTConvolutionLayer::~NEFFTConvolutionLayer() = default;

void NEFFTConvolutionLayer::configure(ITensor                   *input,
                                      const ITensor             *weights,
                                      const ITensor             *biases,
                                      ITensor                   *output,
                                      const PadStrideInfo       &conv_info,
                                      const ActivationLayerInfo &act_info,
                                      bool                       enable_fast_math)
{
    ARM_COMPUTE_TRACE_EVENT(ARM_COMPUTE_PROF_CAT_CPU, ARM_COMPUTE_PROF_LVL_CPU, "NEFFTConvolutionLayer::configure");
    ARM_COMPUTE_UNUSED(enable_fast_math);
    ARM_COMPUTE_LOG_PARAMS(input, weights, biases, output, conv_info, act_info, enable_fast_math);

    _original_weights = weights;
    _original_bias    = biases;

    // Flat if bias addition is required
    _has_bias = biases != nullptr;

    // Get indices for the width and height
    const size_t idx_width = get_data_layout_dimension_index(input->info()->data_layout(), DataLayoutDimension::WIDTH);
    const size_t idx_height =
        get_data_layout_dimension_index(input->info()->data_layout(), DataLayoutDimension::HEIGHT);

    // Input shape, kernel size and output tile
    const Size2D input_dims =
        Size2D(input->info()->tensor_shape()[idx_width], input->info()->tensor_shape()[idx_height]);
    const Size2D kernel_size =
        Size2D(weights->info()->tensor_shape()[idx_width], weights->info()->tensor_shape()[idx_height]);
    const Size2D pad_valid = Size2D(pad_decomposable(input_dims.x() + kernel_size.x() - 1),
                                    pad_decomposable(input_dims.y() + kernel_size.y() - 1));
    // Tensors to use
    ITensor       *input_to_use   = input;
    const ITensor *weights_to_use = weights;
    ITensor       *output_to_use  = _has_bias ? &_bias_output : output;

    // Permute bias
    if (biases != nullptr)
    {
        _permute_bias_func.configure(biases, &_permuted_bias, PermutationVector(1U, 2U, 0U));
        _permuted_bias.info()->set_data_layout(DataLayout::NCHW);
    }

    // Permute input if needed
    _needs_permute = input->info()->data_layout() == DataLayout::NHWC;
    if (_needs_permute)
    {
        _memory_group.manage(&_permuted_input);
        // Configure the function to transform the input tensor from NHWC -> NCHW
        _permute_input_func.configure(input, &_permuted_input, PermutationVector(1U, 2U, 0U));
        _permuted_input.info()->set_data_layout(DataLayout::NCHW);

        // Configure the function to transform the weights tensor from HWI -> IHW
        _permute_weights_func.configure(weights, &_permuted_weights, PermutationVector(1U, 2U, 0U));
        _permuted_weights.info()->set_data_layout(DataLayout::NCHW);

        input_to_use   = &_permuted_input;
        weights_to_use = &_permuted_weights;
    }