EthosnRefPooling2d.hpp

Go to the documentation of this file.

#pragma once
#include "TensorBufferArrayView.hpp"

#include <armnn/Descriptors.hpp>
#include <armnn/Tensor.hpp>

#include "EthosnRefWorkloadUtils.hpp"
#include "EthosnRefConvImpl.hpp"

#include <armnn/Exceptions.hpp>

#include <armnnUtils/DataLayoutIndexed.hpp>

#include <armnn/utility/NumericCast.hpp>

#include <limits>
#include <algorithm>
#include <functional>
#include <cmath>

namespace armnn
{

constexpr std::pair<unsigned int, unsigned int> GetAccumulationLimits(const unsigned int i,
                                                                      const unsigned int inputSize,
                                                                      const unsigned int padBefore,
                                                                      const unsigned int poolSize,
                                                                      const unsigned int stride)
{
    const unsigned int start = std::max(i * stride, padBefore) - padBefore;
    const unsigned int end   = std::min(std::max((i * stride) + poolSize, padBefore) - padBefore, inputSize);

    return { start, end };
}

// Computes the Pooling2d operation.
template<typename InputType>
void EthosnRefPooling2d(const Pooling2dQueueDescriptor& data,
                     const InputType* inputData)
{

    const TensorInfo& inputInfo  = armnn::ethosnref::GetTensorInfo(data.m_Inputs[0]);
    const TensorInfo& outputInfo = armnn::ethosnref::GetTensorInfo(data.m_Outputs[0]);

    TensorBufferArrayView<InputType> output(outputInfo.GetShape(),
                                            armnn::ethosnref::GetOutputTensorData<InputType>(0, data),
                                            data.m_Parameters.m_DataLayout);

    const armnnUtils::DataLayoutIndexed dataLayout(data.m_Parameters.m_DataLayout);
    auto channelsIndex = dataLayout.GetChannelsIndex();
    auto heightIndex = dataLayout.GetHeightIndex();
    auto widthIndex = dataLayout.GetWidthIndex();

    const unsigned int batchSize      = outputInfo.GetShape()[0];
    const unsigned int channels       = outputInfo.GetShape()[channelsIndex];
    const unsigned int outputHeight   = outputInfo.GetShape()[heightIndex];
    const unsigned int outputWidth    = outputInfo.GetShape()[widthIndex];
    const unsigned int inputHeight    = inputInfo.GetShape()[heightIndex];
    const unsigned int inputWidth     = inputInfo.GetShape()[widthIndex];
    const unsigned int paddingLeft    = data.m_Parameters.m_PadLeft;
    const unsigned int paddingTop     = data.m_Parameters.m_PadTop;
    const PaddingMethod method        = data.m_Parameters.m_PaddingMethod;
    const unsigned int strideX        = data.m_Parameters.m_StrideX;
    const unsigned int strideY        = data.m_Parameters.m_StrideY;
    const unsigned int poolHeight     = data.m_Parameters.m_PoolHeight;
    const unsigned int poolWidth      = data.m_Parameters.m_PoolWidth;

    unsigned int inputDataStrides[4] = {};
    if (data.m_Parameters.m_DataLayout == DataLayout::NHWC)
    {
        inputDataStrides[0] = inputHeight * inputWidth  * channels;
        inputDataStrides[1] = 1;
        inputDataStrides[2] = inputWidth * channels;
        inputDataStrides[3] = channels;
    }
    else
    {
        inputDataStrides[0] = inputWidth * inputHeight * channels;
        inputDataStrides[1] = inputWidth * inputHeight;
        inputDataStrides[2] = inputWidth;
        inputDataStrides[3] = 1;
    }

    assert(data.m_Parameters.m_PoolType == armnn::PoolingAlgorithm::Average);

    for (unsigned int n = 0; n < batchSize; ++n)
    {
        for (unsigned int c = 0; c < channels; ++c)
        {
            for (unsigned int yOutput = 0; yOutput < outputHeight; ++yOutput)
            {
                for (unsigned int xOutput = 0; xOutput < outputWidth; ++xOutput)
                {
                    int result = 0;
                    int poolAreaSize = numeric_cast<int>(poolWidth * poolHeight);

                    const auto ySpan = GetAccumulationLimits(yOutput, inputHeight, paddingTop, poolHeight, strideY);
                    const auto xSpan = GetAccumulationLimits(xOutput, inputWidth, paddingLeft, poolWidth, strideX);

                    for (auto yInput = ySpan.first; yInput < ySpan.second; ++yInput)
                    {
                        for (auto xInput = xSpan.first; xInput < xSpan.second; ++xInput)
                        {
                            unsigned int inputIndex = n * inputDataStrides[0] +
                                                      c * inputDataStrides[1] +
                                                      yInput * inputDataStrides[2] +
                                                      xInput * inputDataStrides[3];

                            result += numeric_cast<int>(inputData[inputIndex]);
                        }
                    }

                    if (method == PaddingMethod::Exclude)
                    {
                        poolAreaSize = numeric_cast<int>((xSpan.second - xSpan.first) * (ySpan.second - ySpan.first));
                    }

                    result = (result + (poolAreaSize / 2)) / poolAreaSize;

                    output.Get(n, c, yOutput, xOutput) = numeric_cast<InputType>(result);
                }
            }
        }
    }
}
} //namespace armnn