Method configure

src/cpu/operators/CpuPool2d.cpp:53–95

CpuPool2d::~CpuPool2d() = default;

void CpuPool2d::configure(ITensorInfo *src, ITensorInfo *dst, const PoolingLayerInfo &pool_info, ITensorInfo *indices)
{
    ARM_COMPUTE_TRACE_EVENT(ARM_COMPUTE_PROF_CAT_CPU, ARM_COMPUTE_PROF_LVL_CPU, "CpuPool2d::configure");
    ARM_COMPUTE_LOG_PARAMS(src, dst, pool_info, indices);

    // Check if we can run assembly kernels. Currently, indices are not supported by those kernels
    const bool run_optimised =
        bool(kernels::CpuPool2dAssemblyWrapperKernel::validate(src, dst, pool_info)) && (indices == nullptr);

    // Get data layout
    _data_layout = pool_info.data_layout == DataLayout::UNKNOWN ? src->data_layout() : pool_info.data_layout;

    // Check if we have Global Pooling Layer
    const unsigned int idx_width  = get_data_layout_dimension_index(_data_layout, DataLayoutDimension::WIDTH);
    const unsigned int idx_height = get_data_layout_dimension_index(_data_layout, DataLayoutDimension::HEIGHT);
    _is_global_pooling_layer      = (src->dimension(idx_width) == pool_info.pool_size.width) &&
                               (src->dimension(idx_height) == pool_info.pool_size.height);
    _use_kernel_indices = pool_info.use_kernel_indices;

    if (run_optimised)
    {
        const CPUInfo     &ci          = NEScheduler::get().cpu_info();
        const unsigned int num_threads = NEScheduler::get().num_threads();

        auto pooling_wrapper = std::make_unique<kernels::CpuPool2dAssemblyWrapperKernel>();
        ARM_COMPUTE_ERROR_ON(pooling_wrapper == nullptr);
        pooling_wrapper->configure(src, dst, pool_info, ci);

        // Get kernel's memory requirements
        constexpr size_t alignment      = 4096;
        const size_t     workspace_size = pooling_wrapper->get_working_size(num_threads);
        _aux_mem[0] = MemoryInfo(TensorType::ACL_INT_0, MemoryLifetime::Temporary, workspace_size, alignment);

        _asm_glue = std::move(pooling_wrapper);
    }
    else
    {
        // Configure pooling kernel
        auto k = std::make_unique<kernels::CpuPool2dKernel>();
        k->configure(src, dst, pool_info, indices);
        _pooling_layer_kernel = std::move(k);
    }
}

Status CpuPool2d::validate(const ITensorInfo      *src,
                           const ITensorInfo      *dst,