Method configure

src/cpu/operators/CpuWinogradConv2d.cpp:174–322

CpuWinogradConv2d::~CpuWinogradConv2d() = default;

void CpuWinogradConv2d::configure(const ITensorInfo         *src,
                                  const ITensorInfo         *weights,
                                  const ITensorInfo         *biases,
                                  ITensorInfo               *dst,
                                  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, "CpuWinogradConv2d::configure");
    ARM_COMPUTE_ERROR_ON_NULLPTR(src, weights, dst);
    ARM_COMPUTE_ERROR_THROW_ON(validate(src, weights, biases, dst, conv_info, act_info, enable_fast_math));
    ARM_COMPUTE_LOG_PARAMS(src, weights, biases, dst, conv_info, act_info, enable_fast_math);
    ARM_COMPUTE_UNUSED(biases);
    const DataType data_type = src->data_type();
    uint32_t       nthreads  = NEScheduler::get().num_threads();
    _data_layout             = src->data_layout();
    const Tensor4DShape kernel_shape{internal_get_shape(weights)};

    bool success = get_winograd_kernel_implementation(src, weights, dst, conv_info, act_info, enable_fast_math,
                                                      &_winograd_impl, _conv_args);

    ARM_COMPUTE_EXIT_ON_MSG_VAR(!success, "Unsupported kernel size: %d x %d.\n", kernel_shape.n_rows,
                                kernel_shape.n_cols);
    ARM_COMPUTE_LOG_MSG_WITH_FORMAT_ACL(arm_compute::logging::LogLevel::INFO, "Using input transform: %s\n",
                                        _winograd_impl.input_transform->get_name().c_str());
    ARM_COMPUTE_LOG_MSG_WITH_FORMAT_ACL(arm_compute::logging::LogLevel::INFO, "Using weight transform: %s\n",
                                        _winograd_impl.input_transform->get_name().c_str());
    ARM_COMPUTE_LOG_MSG_WITH_FORMAT_ACL(arm_compute::logging::LogLevel::INFO, "Using output transform: %s\n",
                                        _winograd_impl.input_transform->get_name().c_str());

    const bool has_impl = ((_winograd_impl.input_transform != nullptr) &&
                           (_winograd_impl.output_transform != nullptr) && (_winograd_impl.gemm_args != nullptr));
    if (has_impl)
    {
        // Determine how much working space is required, allocate it.
        const size_t input_workspace_size =
            _winograd_impl.input_transform->get_working_space_size(*_conv_args, nthreads);
        const size_t output_workspace_size =
            _winograd_impl.output_transform->get_working_space_size(*_conv_args, nthreads);

        TensorInfo input_workspace_info(TensorShape(input_workspace_size), 1, DataType::U8);
        TensorInfo output_workspace_info(TensorShape(output_workspace_size), 1, DataType::U8);
        _input_workspace  = input_workspace_info;
        _output_workspace = output_workspace_info;

        const auto &wds = _winograd_impl.winograd_spec;

        // Preparing winograd transformed input tensor
        const size_t     data_type_size    = src->element_size();
        const uint32_t   m                 = _winograd_impl.gemm_args->_Msize; // Total number of tiles
        const uint32_t   k                 = _winograd_impl.gemm_args->_Ksize; // Input channels
        const uint32_t   n                 = _winograd_impl.gemm_args->_Nsize; // Output channels
        const uint32_t   n_gemms           = _winograd_impl.gemm_args->_nmulti;
        const uint32_t   n_batches         = _winograd_impl.gemm_args->_nbatches;
        constexpr size_t storage_alignment = 64;

        const TensorShape a_shape(k, m, n_batches, n_gemms);
        Strides           a_strides(data_type_size);